pytorch导入预训练参数

导入预训练参数

在训练模型的过程中，常需要用到预训练的技巧，即导入在另一个数据集上训练好的参数作为模型参数的初始化，以加快现有模型训练的收敛速度。

1.导入pytorch库中的预训练模型

这些模型属于pytorch内置的模型——Alexnet、VGG、Resnet、DenseNet等，有一个参数pretrained用来控制是否导入预训练参数（这些预训练参数是在ImageNet数据集上训练得到的），代码示例：

// An highlighted block
import torchvision.models as models

model = models.alexnet(pretrained=True)
model = models.vgg19(pretrained=True)
model = models.resnet18(pretrained=True)
model = models.googlenet(pretrained=True)
model = models.densenet121(pretrained=True)
model = models.inception_v3(pretrained=True)

2.有参数文件xxx.pth的情况下，导入参数

其中’./xxx.pth’是文件路径

import torch
import torchvision.models as models

model = MyModle()
model.load_state_dict(torch.load('./xxx.pth'))

3.导入部分参数

这种情况下就是指，模型只需要导入部分预训练参数。比如说，我们的模型（Mymodel）有部分层和vgg19是一样的，然后只想导入这几层的预训练参数。在pytorch中每层模型的参数是以{key:value}的字典形式存储的。key指网络层的名称，value指层的参数。

// An highlighted block
import torchvision.models as models

//自己的模型的初始化并提取网络参数
mymodel = MyModel() 
mymodel_dict = mymodel.state_dict()

//得到vgg19的预训练参数
vgg19 = models.vgg19(pretrained=True)
vgg19_dict = vgg19.state_dict()

//保留vgg19参数中，key（层的名称）存在于mymodel的键值对
state_dict = {k: v for k, v in vgg19_dict.items() if k in mymodel_dict.keys()}

//更新mymodel的参数
mymodel_dict.update(state_dict)

//导入更新后的参数
mymodel.load_state_dict(mymodel_dict)

要注意的点就是，Mymodel中的网络层的名称必须要和vgg19对应层的名称一致，才能导入，不然就会引发keys miss等的错误。

打印一个模型结构的方法：

直接print()。

import torchvision.models as models

model = models.vgg19(pretrained=True)
print(model)

安装torchsummary的包,更为直观，根据输入图像的大小会输出每层的输出大小，以及参数量等。

import torchvision.models as models
from torchsummary import summary

vgg = models.vgg19().cuda()
summary(vgg, (3, 224, 224))