Pytorch保存和加载模型主要有两种方式:
保存:
torch.save(model.state_dict(),"model.pth")
加载:
model.load_state_dict(torch.load("model.pth"))
model.eval()
该方法仅需要保存训练后的模型学习参数。使用torch.save()函数保存模型的state_dict将提供最大的灵活性,以便以后恢复模型,推荐使用此方法来保存模型。常见的PyTorch文件格式是使用.pt
或.pth
文件扩展名保存模型。
注意:
有时候可能由于其他原因模型训练过程中被中断了,这时候就需要保存模型的状态,从终止状态进行训练。
定义模型的状态并保存:
state={
'model':model.state_dict(),'optimizer':optimizer.state_dict(),'epoch':epoch}
torch.save(state,path)
加载继续训练:
checkpoint=torch.load(path)
model.load_state_dict(checkpoint('model'))
optimizer.load_state_dict(checkpoint['optimizer'])
epoch=checkpoint['epoch']
保存:
torch.save(model,"whole_model.pth")
加载:
model=torch.load("whole_model.pth")
model.eval()
这种保存/加载过程使用最直观的语法,并且涉及最少的代码。但以这种方式保存模型将使用Python的pickle模块保存整个模块。这种方法的缺点使序列化的数据保存到特定的类,并且在保存模型时使用确切的目录结构。这样做的原因是因为pickle不会保存模型类本身,而是将其保存到包含类的文件的路径,该路径在加载时使用。因此,在其他项目中使用或重构后,代码可能会以各种方式中断。
File "load_model.py", line 9, in load_model
ckpt = torch.load(weights, map_location=device)
File "/usr/local/python378/lib/python3.7/site-packages/torch/serialization.py", line 594, in load
return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
File "/usr/local/python378/lib/python3.7/site-packages/torch/serialization.py", line 853, in _load
result = unpickler.load()
ModuleNotFoundError: No module named 'models'
训练时采用第二种方式保存整个模型以便于在其他地方调用测试,而该方式保存模型会使序列化的数据保存到特定的类,并且依赖该类文件的特定的目录结构,该路径在加载时使用。因此,在上面项目中调用其他地方保存的模型时由于缺少models路径而找不到models模块。
使用Netron打开保存的整个网络:
上图截取的一部分显示网络层的type
是带有目录结构models.common
的类型,如果加载时目录变了就会导致无法正常导入。
参考github相关解答
def convert_model(model, input=torch.tensor(torch.rand(size=(1,3,112,112)))):
model = torch.jit.trace(self.model, input)
torch.jit.save(model,'/home/projects/models/model.tjm')
然后加载模型:
# load the model
self.model = torch.jit.load('/home/projects/models/model.tjm')
self.model = checkpoint['model'].module
# create the new checkpoint based on what you need
torch.save({
'state_dict' : self.model.state_dict(), 'use_se':True},
'/home/projects/models/best_checkpoint.pth')