Pytorch之保存和加载预训练的模型

在深度学习中会用到迁移学习的方法，也就是我们把在其它数据集上训练比较好的model拿到我们的模型上来进行finetune,这样避免了我们重新去花费时间去训练模型，比如vgg16提取图像特征的这个模型，大大节省了我们训练的时间。这个过程我们就涉及到加载预训练的模型，有的时候我们需要加载整个模型，有时候我们需要模型的一个部分，因此在本文中将会对在Pytroch这个框架中如何加载预训练的模型做以阐述。说到加载预训练好的模型那就不得不说我们如何保存训练好的模型，训练一个model，将这个训练的结果保存起来，然后再测试集上进行测试，这个也是本文将要说的地方。因此就分为两个部分：模型的保存和模型的加载。

1、模型的保存

1）仅保存参数

torch.save(model.state_dict(),path)
#model 模型
#path 保存的路径,保存的文件都是以.pkl或者.pth的形式保存的

2)  既保存模型结构也保存参数。

torch.save(model,path)
#model 模型
#path 保存的路径,保存的文件都是以.pkl或者.pth的形式保存的

 

2、模型的加载