Pytorch模型保存与加载，并在加载的模型基础上继续训练

Posted 2023-03-08

参考技术A

pytorch保存模型非常简单，主要有两种方法：

一般地，采用一条语句即可保存参数：

其中model指定义的模型 实例变量 ，如 model=vgg16( ), path是保存参数的路径，如 path=\'./model.pth\' , path=\'./model.tar\', path=\'./model.pkl\', 保存参数的文件一定要有后缀扩展名。

特别地，如果还想保存某一次训练采用的优化器、epochs等信息，可将这些信息组合起来构成一个字典，然后将字典保存起来：

针对上述第一种情况，也只需要一句即可加载模型：

针对上述第二种以字典形式保存的方法，加载方式如下：

需要注意的是，只保存参数的方法在加载的时候要事先定义好跟原模型一致的模型，并在该模型的实例对象(假设名为model)上进行加载，即在使用上述加载语句前已经有定义了一个和原模型一样的Net, 并且进行了实例化 model=Net( ) 。

另外，如果每一个epoch或每n个epoch都要保存一次参数，可设置不同的path，如 path=\'./model\' + str(epoch) +\'.pth\'，这样，不同epoch的参数就能保存在不同的文件中，选择保存识别率最大的模型参数也一样，只需在保存模型语句前加个if判断语句即可。

下面给出一个具体的例子程序，该程序只保存最新的参数：

在训练模型的时候可能会因为一些问题导致程序中断，或者常常需要观察训练情况的变化来更改学习率等参数，这时候就需要加载中断前保存的模型，并在此基础上继续训练，这时候只需要对上例中的 main() 函数做相应的修改即可，修改后的 main() 函数如下：

以上方法，如果想在命令行进行操作执行，都只需加入argpase模块参数即可，相关方法可参考我的 博客

用法可参照上例。
这篇博客是一个快速上手指南，想深入了解PyTorch保存和加载模型中的相关函数和方法，请移步我的这篇博客： PyTorch模型保存深入理解