pytorch保存模型非常简单,主要有两种方法:
- 只保存参数;(官方推荐)
- 保存整个模型 (结构+参数)。
由于保存整个模型将耗费大量的存储,故官方推荐只保存参数,然后在建好模型的基础上加载。本文介绍两种方法,但只就第一种方法进行举例详解。
一、只保存参数
1.保存
一般地,采用一条语句即可保存参数:
torch.save(model.state_dict(), path)
其中model指定义的模型实例变量,如 model=vgg16( ), path是保存参数的路径,如 path='./model.pth' , path='./model.tar', path='./model.pkl', 保存参数的文件一定要有后缀扩展名。
特别地,如果还想保存某一次训练采用的优化器、epochs等信息,可将这些信息组合起来构成一个字典,然后将字典保存起来:
state = {'model': model.state_dict(), 'optimizer': optimize