用 PyTorch 迁移学习（Transfer Learning）实现图像分类。注意图像数量的分布、图像大小的分布、冻结前面的参数、统计模型参数数量。

用 PyTorch 迁移学习（Transfer Learning）实现图像分类

作者: PyTorch 中文网 

https://github.com/WillKoehrsen/pytorch_challenge/blob/master/Transfer%20Learning%20in%20PyTorch.ipynb

 

教程将手把手教你用 PyTorch 实现迁移学习（Transfer Learning）来做图像分类。数据库我们采用的是 Caltech 101 dataset，这个数据集包含 101 个图像分类，大多数分类只包含 50 张左右的图像，这对于神经网络来讲是远远不够的。那我们就用一个实现训练好的图像分类模型加迁移学习的方法，来实现在这个数据集上的训练。

目录

用 PyTorch 迁移学习（Transfer Learning）实现图像分类

什么是迁移学习

迁移学习的基本思想

准备数据集

图像数量的分布

图像大小的分布

图像增广（Data augmentation）

图像预处理

接下来定义 dataset 和 DataLoader。

ImageNet 的预训练模型

首先加载预训练模型：

冻结前面的参数

统计模型参数数量：

损失函数和优化器

训练

预测

---

什么是迁移学习

迁移学习（Transfer Learning）的基本概念就是当可用的数据集特别少时，从头开始训练一个神经网络往往不会得到很好的结果，于是就从一个预训练模型开始训练，让网络本身已经具备一定的训练基础，然后用小数据集进行微调，便可以得到一个不错的结果。

通常加载预训练模型后，我们冻结模型的部分参数，一般只训练模型的最后几层，这样可以保留整个模型前面对物体特征提取的能力。预训练模型一定要与新的数据集有共同点，比如都是图像分类问题，这行才能有效地把预训练模型里的特征提取能力迁移到新的模型上。

迁移学习的基本思想

下面是迁移学习用于物体识别时的一般过程：

1. 加载预训练模型
2. 冻结模型前面部分的参数
3. 添加可训练的自定义的分类层，或使用原模型的分类层（如果可重用的话）
4. 在新数据集上训练

准备数据集

下载数据集后，我们按照 50%，25%，25% 的比例划分 training，validation，和 testing。目录整理如下：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16

/datadir     /train         /class1         /class2         .         .     /valid         /class1         /class2         .         .     /test         /class1         /class2         .         .

我们来看一下每一个分类里都有多少张图片：

图像数量的分布

可以看到有很多分类的图像都很少，为了达到最好的训练效果，我们后面会用 data augmentation 来增加图像的数量。

我们再来看一下图像大小的分布：

图像大小的分布

我们要用的预训练模型是基于 ImageNet 的，它的训练图像大小是 224 x 224，所以我们还需要对我们的数据集的图像进行大小缩放。

图像增广（Data augmentation）

图像增广一般用来人工产生不同的图像，比如对图像进行旋转、翻转、随机裁剪、缩放等等。这里我们选择在训练阶段对输入进行增广，比如说我们训练了 20 个 epoch，那么每个 epoch 里网络看到的输入图像都会略微不同。

图像预处理

在 PyTorch 里，我们用 transforms 进行图像预处理。首先我们定义 training 和 validation 的预处理方式：

 

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25

from torchvision import transforms   # Image transformations image_transforms = {     # Train uses data augmentation     'train':     transforms.Compose([         transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),         transforms.RandomRotation(degrees=15),         transforms.ColorJitter(),         transforms.RandomHorizontalFlip(),         transforms.CenterCrop(size=224),  # Image net standards         transforms.ToTensor(),         transforms.Normalize([0.485, 0.456, 0.406],                              [0.229, 0.224, 0.225])  # Imagenet standards     ]),     # Validation does not use augmentation     'valid':     transforms.Compose([         transforms.Resize(size=256),         transforms.CenterCrop(size=224),         transforms.ToTensor(),         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])     ]), }

接下来定义 dataset 和 DataLoader。

用 datasets.ImageFolder 来定义 dataset 时 PyTorch 可以自动将图片与对应的文件夹分类对应起来，而且应用我们上面定义好的 transformers，然后 dataset传入到 DataLoader 里，DataLoader 在每一个循环会自动生成 batchsize 大小的图像和 label。

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16

from torchvision import datasets from torch.utils.data import DataLoader   # Datasets from folders data = {     'train':     datasets.ImageFolder(root=traindir, transform=image_transforms['train']),     'valid':     datasets.ImageFolder(root=validdir, transform=image_transforms['valid']), }   # Dataloader iterators, make sure to shuffle dataloaders = {     'train': DataLoader(data['train'], batch_size=batch_size, shuffle=True),     'val': DataLoader(data['valid'], batch_size=batch_size, shuffle=True) }

我们可以看一下 DataLoader 的输出：

1 2 3 4

trainiter = iter(dataloaders['train']) features, labels = next(trainiter) features.shape, labels.shape (torch.Size([128, 3, 224, 224]), torch.Size([128]))

ImageNet 的预训练模型

PyTorch 自带了很多 ImageNet 上的预训练模型，详细列表见这里。下表是各个模型的性能对比：

本教程将选用 VGG-16 的预训练模型。

首先加载预训练模型：

1 2	from torchvision import models model = model.vgg16(pretrained=True)

冻结前面的参数

我们只训练这个模型最后的全链接层，所以首先我们要冻结前面的参数：

1 2	for param in model.parameters():     param.requires_grad = False

因为我们的数据集只有 100 个分类，所以要在模型最后面加上几层使得模型的最终输出跟我们的类别数目一样：

1 2 3 4 5 6 7 8

import torch.nn as nn # Add on classifier model.classifier[6] = nn.Sequential(                       nn.Linear(n_inputs, 256),                       nn.ReLU(),                       nn.Dropout(0.4),                       nn.Linear(256, n_classes),                                         nn.LogSoftmax(dim=1))

我们来看一下整个模型：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16

model.classifier Sequential(   (0): Linear(in_features=25088, out_features=4096, bias=True)   (1): ReLU(inplace)   (2): Dropout(p=0.5)   (3): Linear(in_features=4096, out_features=4096, bias=True)   (4): ReLU(inplace)   (5): Dropout(p=0.5)   (6): Sequential(     (0): Linear(in_features=4096, out_features=256, bias=True)     (1): ReLU()     (2): Dropout(p=0.4)     (3): Linear(in_features=256, out_features=100, bias=True)     (4): LogSoftmax()   ) )

统计模型参数数量：

1 2 3 4 5 6 7

total_params = sum(p.numel() for p in model.parameters()) print(f'{total_params:,} 参数总数.') total_trainable_params = sum(     p.numel() for p in model.parameters() if p.requires_grad) print(f'{total_trainable_params:,} 可训练参数总数.') 135,335,076 参数总数. 1,074,532 可训练参数总数.

损失函数和优化器

这里我们使用的损失函数是 negative log likelihood (NLL)，优化器是 Adam。

 

1 2 3 4

from torch import optim # Loss and optimizer criteration = nn.NLLLoss() optimizer = optim.Adam(model.parameters())

训练

下面是训练的伪代码，大家理解一下其中的思想：

1 2 3 4 5 6 7 8 9 10 11

# 伪代码 for epoch in range(n_epochs):   for data, targets in trainloader:     # Generate predictions     out = model(data)     # Calculate loss     loss = criterion(out, targets)     # Backpropagation     loss.backward()     # Update model parameters     optimizer.step()

完整的训练代码大家可以参考完整代码。下图是训练的 loss 和 accuracy 曲线：

预测

模型训练好后，就可以做预测了：

1 2 3 4 5 6

for data, targets in testloader:     log_ps = model(data)     # Convert to probabilities     ps = torch.exp(log_ps) ps.shape() (128, 100)

因为我们对所有分类都有输出概率，所以我们要找出概率最大的那个类别来作为最后的预测值：

1 2 3 4 5

# Find predictions and correct pred = torch.max(ps, dim=1) equals = pred == targets # Calculate accuracy accuracy = torch.mean(equals)

我们来看一些输出：

可以看到，模型的预测还是很准确的。

这是一个基本的用 PyTorch 实现迁移学习的训练过程，完整的代码见这个链接。