python-pytorch 关于torch.load和torch.load_state_dict
最近在使用pytorch训练和加载模型时遇到了一些很玄学的问题,研究了一下,总结如下:
1、关于模型保存和加载
1、如果保存时使用了torch.save( model, “model.pth”),那么模型加载时只能使用torch.load(),而不能使用torch.load_state_dict(torch.load(model_path)),如果强行使用torch.load_state_dict(torch.load(model_path)),会报Net对象没有copy属性。
import torch
torch.save(model, 'model.pth')
model = torch.load('model.pth')
2、如果保存时使用了torch.save( model.state_dict(),“model.pth”),那么模型加载时只能使用model.load_state_dict(torch.Load(“model.pth”))。
import torch
torch.save(model.state_dict(), 'model.pth')
state_dict = torch.load('model.pth')
model.load_state_dict(state_dict)
3、两种方式在模型加载时,都需要引入或定义模型结构,且模型结构的名称需要与训练时一致,不能使用别名。
训练时:class Net,测试时:定义或加载仍然是class Net,而不能from ··· import Net as resNet。
# 两个Net的结构不同
from model1 import Net
from model2 import Net
2、关于加载模型结构
关于在一个脚本中同时加载两个模型,在同时导入两个不同结构的Net的情况下,即使Net是同名,但加载的模型会自动找到自己训练时对应的那个结构。