PyTorch系列 | 快速入門遷移學習

原題 | TRANSFER LEARNING TUTORIAL

做者 | Sasank Chilamkurthy

原文 | pytorch.org/tutorials/b…

譯者 | kbsc13("算法猿的成長"公衆號做者)

聲明 | 翻譯是出於交流學習的目的，歡迎轉載，但請保留本文出於，請勿用做商業或者非法用途

簡介

本次教程主要介紹如何用深度學習實現遷移學習。更多更詳細的遷移學習知識能夠查看 cs231n 課程--cs231n.github.io/transfer-le…

實際應用中，不多人會從頭開始，經過隨機初始化來訓練一個卷積神經網絡，由於擁有足夠數量的數據集太少了。一般，你們都會選擇一個在比較大的數據集（好比 ImageNet 數據集，1000個類別總共120萬張圖片）上訓練好的預訓練模型，而後對卷積神經網絡進行初始化，或者用於提取特徵。

遷移學習的兩大主要應用場景：

• 微調網絡：採用預訓練模型對網絡進行初始化，而不是隨機初始化，這種作法能夠更快收斂，同時也能取得更好的效果。
• 做爲一個特徵提取器：這種應用會將除了最後的全鏈接層外的網絡層都固定權重參數，而後最後的全鏈接層會根據數據集類別數量進行修改輸出，並隨機初始化其權值，而後訓練該層。

本文的教程，代碼中須要導入的模型以下所示：

# License: BSD
# Author: Sasank Chilamkurthy

from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy

plt.ion()   # interactive mode
複製代碼

加載數據

加載數據這部分將採用 torchvision 和 torch.utils.data 兩個模塊。

本次教程的目標是訓練一個二分類模型，類別是螞蟻和蜜蜂，所以數據集中分別包含了 120 張螞蟻和蜜蜂的訓練圖片，而後每一個類別還包含 75 張圖片做爲驗證集。也就是說這個數據集總管只有 390 張圖片，不到一千張圖片，是一個很是小的數據集，若是從頭訓練模型，很難得到很好的泛化能力。所以，本文將對這個數據集採用遷移學習的方法來獲得更好的泛化能力。

獲取本文數據集和代碼，能夠在公衆號後臺回覆「pytorch遷移學習」獲取。

加載數據的代碼以下所示：

# 數據加強方法，訓練集會實現隨機裁剪和水平翻轉，而後進行歸一化
# 驗證集僅僅是裁剪和歸一化，並不會作數據加強
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}
# 數據集所在文件夾
data_dir = 'data/hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
複製代碼

可視化圖片

首先可視化一些訓練圖片，以便於更好理解數據加強。代碼以下所示：

# 圖片展現的函數
def imshow(inp, title=None):
    """Imshow for Tensor."""
    # 逆轉操做，從 tensor 變回 numpy 數組須要轉換通道位置
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    # 從歸一化後變回原始圖片
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # pause a bit so that plots are updated


# 獲取一個 batch 的訓練數據
inputs, classes = next(iter(dataloaders['train']))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_names[x] for x in classes])
複製代碼

展現的圖片以下所示：

訓練模型

加載數據後，就是開始進行訓練模型，這裏會介紹如下兩個內容：

• 制定學習率的策略
• 保存最佳模型

在下面的代碼中，參數 scheduler 是採用 torch.optim.lr_scheduler 初始化的 LR 策略對象：

# 訓練模型的函數
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # 每一個 epoch 都分爲訓練階段和驗證階段
        for phase in ['train', 'val']:
            # 注意訓練和驗證階段，須要分別對 model 的設置
            if phase == 'train':
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 清空參數的梯度
                optimizer.zero_grad()

                # 只有訓練階段才追蹤歷史
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # 訓練階段才進行反向傳播和參數的更新
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 記錄 loss 和 準確率
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # 載入最好的模型參數
    model.load_state_dict(best_model_wts)
    return model
複製代碼

上述函數實現了模型的訓練，在一個 epoch 中分爲訓練和驗證階段，訓練階段天然須要前向計算加反向傳播，並更新網絡層的參數，但驗證階段只須要前向計算，而後記錄 loss 和驗證集上的準確率便可。

此外就是須要設置保存模型的條件，這裏是當每次驗證集的準確率都高於以前最好的準確率時，保存模型。

可視化模型的預測結果

下面定義了一個可視化模型預測結果的函數，用於展現圖片和模型對該圖片的預測類別信息：

# 可視化模型預測結果，即展現圖片和模型對該圖片的預測類別信息，默認展現 6 張圖片
def visualize_model(model, num_images=6):
    was_training = model.training
    model.eval()
    images_so_far = 0
    fig = plt.figure()

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            for j in range(inputs.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images//2, 2, images_so_far)
                ax.axis('off')
                ax.set_title('predicted: {}'.format(class_names[preds[j]]))
                imshow(inputs.cpu().data[j])

                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return
        model.train(mode=was_training)
複製代碼

微調網絡

這部分就是本次遷移學習的核心內容，前面都是正常的加載數據、定義訓練過程的代碼，這裏介紹的就是如何進行微調網絡，代碼以下所示：

# 加載 resnet18 網絡模型，而且設置加載預訓練模型
model_ft = models.resnet18(pretrained=True)
num_ftrs = model_ft.fc.in_features
# 修改輸出層的輸出數量，本次採用的數據集類別爲 2
model_ft.fc = nn.Linear(num_ftrs, 2)

model_ft = model_ft.to(device)

criterion = nn.CrossEntropyLoss()

# 對全部網絡層參數進行更新
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# 學習率策略，每 7 個 epochs 乘以 0.1
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
複製代碼

這一步由於是設置加載預訓練模型，因此運行後會下載預訓練的 resnet18 網絡模型文件

訓練和驗證

接下來開始正式訓練網絡模型了，代碼以下所示，若是採用 cpu，大約須要 15-25 分鐘的過程，而若是採用 gpu，那麼速度就很快了，基本一分鐘左右就訓練完成了。

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=25)
複製代碼

訓練結果：

可視化模型預測結果：

visualize_model(model_ft)
複製代碼

可視化結果以下：

用於特徵提取器

剛剛是用於微調網絡，即將預訓練模型用於初始化網絡層的參數，接下來介紹遷移學習的第二種用法，做爲特徵提取器，也就是固定預訓練模型部分的網絡層的權值參數，這部分的實現代碼，以下所示，其中須要將卷積層部分的參數固定，即設置 requires_grad==False ，這樣在反向傳播過程就不會計算它們的梯度，更多內存能夠查看 pytorch.org/docs/notes/…

model_conv = torchvision.models.resnet18(pretrained=True)
# 固定卷積層的權重參數
for param in model_conv.parameters():
    param.requires_grad = False

# 新的網絡層的參數默認 requires_grad=True
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)

model_conv = model_conv.to(device)

criterion = nn.CrossEntropyLoss()

# 只對輸出層的參數進行更新
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)

# 學習率策略，每 7 個 epochs 乘以 0.1
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)
複製代碼

再次進行訓練：

model_conv = train_model(model_conv, criterion, optimizer_conv,
                         exp_lr_scheduler, num_epochs=25)
複製代碼

訓練結果：

可視化網絡的預測結果

visualize_model(model_conv)

plt.ioff()
plt.show()
複製代碼

輸出結果：

小結

本文教程簡單介紹了遷移學習的內容，經過訓練一個二分類模型，簡單介紹遷移學習的兩個用法，微調網絡和用於固定的特徵提取器。

本文代碼地址：

github.com/ccc013/Deep…

獲取本文的代碼和數據集方法：

1. 關注公衆號「算法猿的成長」
2. 公衆號會話界面回覆「pytorch遷移學習"

歡迎關注個人微信公衆號--算法猿的成長，或者掃描下方的二維碼，你們一塊兒交流，學習和進步！