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'时，它能够正常工作。我应该如何解决这个问题？感谢您的回答。 :)

- Hyelin

Answer 1

看起来你在保存时使用了 cuda 状态，现在尝试在 cpu 上使用它或者反过来。为了避免这个错误，一个简单的方法是在加载时传递 map_location 参数。

只需在 torch.load 中传递 map_location=<你想要使用的设备>，就可以正常工作了。另外，请参见 https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-across-devices。

- Umang Gupta

---

state = torch.load('drive/My Drive/MODEL/4 CBAM classifier55',map_location=device) 我加了map_location参数（device是cuda），但是还是不行。出现了相同的错误信息。:( 有没有办法检查我是在哪里保存模型的（CPU还是GPU）？ - Hyelin

我在声明优化器之前改变了“model.to(device)”的位置，现在它可以正常工作了！感谢您的帮助。 - Hyelin

嗯！我看到了，很高兴它起作用了。通常，首先将模型移动到设备上，然后再声明优化器是一个好主意。这可以避免优化器在 CPU 和 GPU 上的某些部分混淆。 - Umang Gupta