pytorch load_state_dict不生效

有时候我们会发现，load一个保存好的模型dict到一个新的模型对象的时候，并没有加载dict里面的数据。

一个简单例子

比如下图中，模型TestM里面定义了一个参数D，如果我不加：.cuda()，那么下面的例子是可以成功加载的，也就是说，print出来的 Loaded和Origin是相等的。

但是，如果对参数D加了 .cuda()，那么，Loaded就和Init 是一样的，并不等于Origin。

结论

模型参数定义不要加.cuda()，应该统一在外面调用model.cuda().

大家可以复制下面代码进行简单测试即可。

import torch
from torch import nn

class TestM(nn.Module):
    def __init__(self):
        super(TestM, self).__init__()
        self.D = nn.Parameter(torch.randn(2, 2).float())  # 成功
        # self.D = nn.Parameter(torch.randn(2, 2).float()).cuda() # 失败
        
    def forward(self, x):
        return x

tm1 = TestM()
print("Origin: ", tm1.D)
torch.save(tm.state_dict(), 'tm_test.pth')

tm2 = TestM()
print("Init: ", tm2.D)
tm2.load_state_dict(torch.load('tm_test.pth'))
print("Loaded: ", tm2.D)