GPU训练完之后转移到CPU测试

1 保存模型

我们保存模型的方式通常有两种方式

1 torch.save(model,r'model.pkl')
2 torch.save(model.state_dict(),r'model.pkl')   

推荐第二种，在这次示例当中，我们采用第二种方式，保存文件。

2 开始操作

1 首先看一下，原本的model 字典 key

for param_tensor in model.state_dict():        # param_tensor 为参数名称
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

2 加载模型

model = My_model()

checkpoint = torch.load(r'model_best4.pkl',map_location='cpu')

查看此时加载后的模型 key

要比我们原本的模型多了module. 总共7位，所以我们需要一个新的字典重新加载模型参数。

new_state_dict = OrderedDict()
for k,v in checkpoint.items() :
    name = k[7:]
    new_state_dict[name] = v

新的字典构造后 ，我们看一下构造对没。

for k,v in new_state_dict.items():
    print(k)

完全没问题！现在就可以把这个字典保存，再加载了。

torch.save(new_state_dict,'new.pkl')
model.load_state_dict(torch.load(r'new.pkl',map_location='cpu'))

问题：
1 为什么我们不能把新建好的字典直接给模型加载，还得保存再加载？
答：
因为使用 model.load_state_dicr(torch.load(r’new.pkl’,map_location=‘cpu’))
只能加载二进制文件，我们保存模型参数的时候经常使用 .th 和 .pkl 后缀就是二进制文件。所以我们需要把新构成的字典转换成二进制文件，最后再加载。