Pytorch：迁移学习时使用部分预训练模型

应用场景

通常在一些现实任务中，我们所能拿到的数据比较少，而深度学习模型的参数更新通常需要大量的数据来驱动，所以我们就可以对某些层用到预训练模型。

如何下载常用的预训练模型

通常在Pytorch的github上可以找到torchvision一些模型源码（比如VGG,ResNet,GoogLeNet等等）：
https://github.com/pytorch/vision/tree/master/torchvision/models

找到源码中的下载地址，在浏览器中下载预训练模型文件。例如vgg.py中就包含以下可使用的模型：

在下载好预训练模型文件之后，我们就可以正式开始了。

1. 按照源码设置相应的参数：

以VGG为例，在官方的模型中：

		self.features = features
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),  # num_classes=1000
        )

主要包括features，avgpool，还有classifier几个模块。features就是卷积模块用来提取特征，avgpool是自适应的平均池化层，classifier是包含全连接层的分类器。比如我现在只需要对80个类建立分类模型，那么我除了最后一个Linear的参数,其他都需要用到预训练模型参数。

• 我的模型设置：

		self.features = nn.Sequential(
            # Layer 1
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, stride=2),
            # Layer 2
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, stride=2),
            # Layer 3
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, stride=2),
            # Layer 4
            nn.Conv2d(256, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, stride=2),
            # Layer 5
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, stride=2)
        )
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512*7*7, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
        )
        # 因为前面可以用预训练模型参数，所以单独把最后一层提取出来
        self.classifier2 = nn.Linear(4096, num_classes)

2. 复制所需部分的模型参数

# 使用部分预训练模型方法1
model = VGG19(num_classes=2, init_weights=False)
model_dict = model.state_dict()
# 加载预训练模型
state_dict = torch.load('pretrained/vgg19-dcbb9e9d.pth')
# 把相同参数名的参数找出来
new_state_dict = {k: v for k, v in state_dict.items() if k in model_dict}
# 更新参数
model_dict.update(new_state_dict)
# 加载模型参数
model.load_state_dict(model_dict)

# 使用部分预训练模型方法2
model = VGG19(num_classes=2, init_weights=False)
# 加载模型参数时，令strict为False
# 表明不严格要求state_dict中的键与该模块中state_dict()函数返回的键匹配
model.load_state_dict(torch.load('pretrained/vgg19-dcbb9e9d.pth'), strict=False)

两种方式加载的预训练模型都是一样的

3. 打印参数验证

for name, para in model.named_parameters():
    print(name, torch.min(para)) # 只打印最小值

可以看到，参数已经成功加载。