问题:

Traceback (most recent call last):
  File "2_generate_PM_CAM_bcss.py", line 47, in <module>
    model.load_state_dict(torch.load(args.weights), strict=False)
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/serialization.py", line 712, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/serialization.py", line 1046, in _load
    result = unpickler.load()
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/serialization.py", line 1016, in persistent_load
    load_tensor(dtype, nbytes, key, _maybe_decode_ascii(location))
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/serialization.py", line 1001, in load_tensor
    wrap_storage=restore_location(storage, location),
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/serialization.py", line 176, in default_restore_location
    result = fn(storage, location)
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/serialization.py", line 158, in _cuda_deserialize
    return obj.cuda(device)
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/_utils.py", line 79, in _cuda
    return new_type(self.size()).copy_(self, non_blocking)
  File "/home/pengzhang/anaconda3/envs/TPRO/lib/python3.8/site-packages/torch/cuda/__init__.py", line 661, in _lazy_new
    return super(_CudaBase, cls).__new__(cls, *args, **kwargs)
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

我在代码中已经设置了:

torch.cuda.set_device(6)

即:我已经设置了在cuda:6上面运行,但是一直显示的是程序运行在cuda:1上面。

解决措施:

state_dict = torch.load(args.weights, map_location=torch.device('cuda:6' if torch.cuda.is_available() else 'cpu'))
model.load_state_dict(state_dict, strict=False)

即通过 map_location参数将模型参数映射到一个现有的设备上。

个人推测:可能是在保存模型权重的时候,将训练显卡型号也保存下来了,这样一来,在加载权重的时候自然就会加载到指定的显卡上面。(可能不是很对,先去吃饭去了,回来再看看,debug了一上午)