一、pth 文件中保存的是什么

import torch
state_dict = torch.load("resnet18.pth")
print(type(state_dict))

---------------
<class 'collections.OrderedDict'>

如上打印输出所示,pth 文件通过 有序字典 来保存模型参数。有序字典与常规字典一样,但是在排序操作方面有一些额外的功能。常规的 dict 是无序的,OrderedDict 能够比 dict 更好地处理频繁的重新排序操作。


二、使用 torch.save()/torch.load() 保存/恢复模型

1. torch.save() 保存模型

  1. 先建立一个字典,保存三个参数。
  2. 调用torch.save(),即可保存对应的pth文件。
state_dict = {‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
torch.save(state_dict , dir)

--------------------------------
torch.save(model.state_dict,dir)

需要注意的是若模型是由 nn.Moudle 类继承的模型,保存 pth 文件时,state_dict 参数需要由model.state_dict 指定。

2. torch.load() 恢复模型

当你想恢复某一阶段的训练(或者进行测试)时,那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(dir)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch'] + 1

参考链接

  1. pytorch模型文件pth详解
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐