训练加载模型gpu映射问题
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.
遇到的错误是因为PyTorch尝试加载一个在不存在的CUDA设备上保存的模型。具体来说,这个模型是在CUDA设备1上保存的,但你当前的环境中只有一个CUDA设备(设备0)
要解决这个问题,需要在torch.load
函数中使用map_location
参数将加载的模型映射到一个有效的CUDA设备(或者映射到CPU,如果你更愿意这样做)。以下是修改代码的两种方式:
1.映射到现有的CUDA设备(如果你想使用GPU):
best_model_path = path + '/' + 'checkpoint.pth'
self.model.load_state_dict(torch.load(best_model_path))
解决方法:上述代码中加gpu映射
best_model_path = path + '/' + 'checkpoint.pth'
map_location = lambda storage, loc: storage.cuda(self.args.gpu)
state_dict = torch.load(best_model_path, map_location=map_location)
2.映射到CPU(如果你没有GPU或者更愿意使用CPU):
import torch
# 加载模型并映射到CPU
best_model_path = 'path_to_your_model_file.pth'
self.model.load_state_dict(torch.load(best_model_path, map_location=torch.device('cpu')))