floodsung / LearningToCompare_FSL

PyTorch code for CVPR 2018 paper: Learning to Compare: Relation Network for Few-Shot Learning (Few-Shot Learning part)
MIT License
1.04k stars 266 forks source link

RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device. #33

Open slivara opened 4 years ago

slivara commented 4 years ago

(base) C:\Users\B4-410\lpthw\LearningToCompare_FSL-master\miniimagenet>python miniimagenet_train_one_shot.py -w 5 -s 1 -b 15 init data folders init neural networks Traceback (most recent call last): File "miniimagenet_train_one_shot.py", line 269, in main() File "miniimagenet_train_one_shot.py", line 150, in main feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_featureencoder" + str(CLASSNUM) +"way" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl"))) File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 426, in load return _load(f, map_location, pickle_module, **pickle_load_args) File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 613, in _load result = unpickler.load() File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 576, in persistent_load deserialized_objects[root_key] = restore_location(obj, location) File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 155, in default_restore_location result = fn(storage, location) File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 131, in _cuda_deserialize device = validate_cuda_device(location) File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 125, in validate_cuda_device device, torch.cuda.device_count())) RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

githubwys commented 4 years ago

i met the same problem.

you have to assign a certain gpu.

find the file miniimagenet_train_one_shot.py from :

feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_featureencoder" + str(CLASSNUM) +"way" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl")))

to : feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_featureencoder" + str(CLASSNUM) +"way" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl"), map_location='cuda:0'))

also the same as the relation_network part.

hope helpful to you