当优化器加载状态字典时 Pytorch / 设备问题（CPU，GPU）

Question

大家好，我是一名学生，从去年夏天开始学习 pytorch。

state = torch.load('drive/My Drive/MODEL/4 CBAM classifier55')

model = MyResNet()
model.load_state_dict(state['state_dict'])

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.0003,betas=(0.5,0.999))
optimizer.load_state_dict(state['optimizer'])

model.to(device)

我写了像上面这样的代码。

RuntimeError                              Traceback (most recent call last)
<ipython-input-26-507493db387a> in <module>()
     56     new_loss.backward()
     57 
---> 58     optimizer.step()
     59 
     60     running_loss += loss.item()

/usr/local/lib/python3.6/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     13         def decorate_context(*args, **kwargs):
     14             with self:
---> 15                 return func(*args, **kwargs)
     16         return decorate_context
     17 

/usr/local/lib/python3.6/dist-packages/torch/optim/adam.py in step(self, closure)
     97 
     98                 # Decay the first and second moment running average coefficient
---> 99                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
    100                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
    101                 if amsgrad:

RuntimeError: expected device cpu but got device cuda:0

当我实现训练代码时，我遇到了这种错误。当我注释掉“optimizer.load_state_dict”时，它运行良好。我怎么解决这个问题？谢谢您的回答。:)

Answer 1

似乎在您保存时state打开cuda，现在尝试使用它cpu，反之亦然。为了避免这个错误，一个简单的方法是将map_location参数传递给load。

只要通过map_location=<device you want to use>就torch.load可以正常工作。另请参阅https://pytorch.org/tutorials/beginner/ saving_loading_models.html# saving-loading-model-across-devices