训练加载模型gpu映射问题

训练加载模型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')))