用 PyTorch 遷移學習（Transfer Learning）實現圖像分類
作者: PyTorch 中文網
教程將手把手教你用 PyTorch 實現遷移學習（Transfer Learning）來做圖像分類。數據庫我們採用的是 Caltech 101 dataset，這個數據集包含 101 個圖像分類，大多數分類只包含 50 張左右的圖像，這對於神經網絡來講是遠遠不夠的。那我們就用一個實現訓練好的圖像分類模型加遷移學習的方法，來實現在這個數據集上的訓練。
目錄
用 PyTorch 遷移學習（Transfer Learning）實現圖像分類
什麼是遷移學習
遷移學習（Transfer Learning）的基本概念就是當可用的數據集特別少時，從頭開始訓練一個神經網絡往往不會得到很好的結果，於是就從一個預訓練模型開始訓練，讓網絡本身已經具備一定的訓練基礎，然後用小數據集進行微調，便可以得到一個不錯的結果。
通常加載預訓練模型後，我們凍結模型的部分參數，一般只訓練模型的最後幾層，這樣可以保留整個模型前面對物體特徵提取的能力。預訓練模型一定要與新的數據集有共同點，比如都是圖像分類問題，這行纔能有效地把預訓練模型裏的特徵提取能力遷移到新的模型上。
遷移學習的基本思想
下面是遷移學習用於物體識別時的一般過程：
- 加載預訓練模型
- 凍結模型前面部分的參數
- 添加可訓練的自定義的分類層，或使用原模型的分類層（如果可重用的話）
- 在新數據集上訓練
準備數據集
下載數據集後，我們按照 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 實現遷移學習的訓練過程，完整的代碼見這個鏈接。