Pytorch迁移学习实战：基于vgg11做MNIST数据集的分类

在对一个新的数据集进行分类时，经常使用一些预训练好的模型，如AlexNet，VGG，Resnet等，并且将网络的最后几层fc进行修改，以适应新数据集的类别个数。

这样做的可行性在于，深度神经网络通过预训练得到数据分层级的特征，比较shallow的层学习到的是一些低层级的语义特征，如边缘信息、颜色、纹理等等，这些特征在不同的分类任务中大多是不变的；最后几层学习到的是一些高层的特征，这样是不同数据集的区别所在，因此通常对VGG，Resnet的最后几层进行修改，利用预训练好的模型进行微调，以在新的数据集上进行分类。

本文利用pytorch中自带的VGG11模型进行迁移学习，使得模型在MNIST数据集上进行分类。首先可以打印看一下vgg11模型的具体结构，模型概述如下：

VGG(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(4): ReLU(inplace=True)
(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(6): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(7): ReLU(inplace=True)
(8): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(9): ReLU(inplace=True)
(10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(11): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(12): ReLU(inplace=True)
(13): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(14): ReLU(inplace=True)
(15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(16): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(17): ReLU(inplace=True)
(18): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(19): ReLU(inplace=True)
(20): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
(classifier): Sequential(
(0): Linear(in_features=25088, out_features=4096, bias=True)
(1): ReLU(inplace=True)
(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=4096, out_features=4096, bias=True)
(4): ReLU(inplace=True)
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=4096, out_features=1000, bias=True)
)
)

可见vgg11最后的输出是1000维向量，如果要满足MNIST数据集，需要将其改为10维的输出。目前我了解到的模型改法有两种：其一是对VGG某一层进行修改，如改变Linear的维度，或Conv2d的通道数；第二种是在VGG的基础上进行layer的添加，如在全连接层前面再加一个Conv+Relu+MaxPool的block。

一、在layer的基础上进行修改

这种做法相对简单。如果仅仅是对已有的layer进行改造，比如将最后的全连接层的神经元个数修改一下，并且最终得到的输出为10维，训练代码如下（下面的代码中将calssifier的(0),(3),(6)三个Linear层的神经元个数进行修改，并修改了输入的卷积层维度）：

import torch
from torchvision import models
from torchvision import datasets
from torchvision import transforms
import torch.nn as nn
from torch.optim import Adam

epoch = 1
lr = 0.0001
batch_size = 64

def build_model():
    vgg11 = models.vgg11(pretrained=True)
    #MNIST为单通道，而原模型输入为3通道，这里仅对卷积层输入维度做出改变
    vgg11.features[0] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    in_channel = 25088
    out_channel = [4096, 1024, 10]
    #修改最后三个fc层
    for i in range(7):
        if i % 3 == 0:
            vgg11.classifier[i] = nn.Linear(in_channel, out_channel[int(i/3)])
            in_channel = out_channel[int(i/3)]

    return vgg11

model = build_model()

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

train_data = datasets.MNIST(root='./data/', train=True, transform=transform, download=True)
test_data = datasets.MNIST(root='./data/', train=<