[PyTorch入門]之遷移學習

遷移學習教程

來自這里。

在本教程中，你將學習如何使用遷移學習來訓練你的網絡。在cs231n notes你可以了解更多關於遷移學習的知識。

    在實踐中，很少有人從頭開始訓練整個卷積網絡(使用隨機初始化)，因為擁有足夠大小的數據集相對較少。相反，通常在非常大的數據集(例如ImageNet，它包含120萬幅、1000個類別的圖像)上對ConvNet進行預訓練，然后使用ConvNet作為初始化或固定的特征提取器來執行感興趣的任務。

兩個主要的遷移學習的場景如下：

• Finetuning the convert：與隨機初始化不同，我們使用一個預訓練的網絡初始化網絡，就像在imagenet 1000 dataset上訓練的網絡一樣。其余的訓練看起來和往常一樣。
• ConvNet as fixed feature extractor：在這里，我們將凍結所有網絡的權重，除了最后的全連接層。最后一個全連接層被替換為一個具有隨機權重的新層，並且只訓練這一層。

#!/usr/bin/env python3

# License: BSD
# Author: Sasank Chilamkurthy

from __future__ import print_function,division

import torch
import torch.nn as nn
import torch.optim as optim
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()   # 交互模式

導入數據

我們使用torchvision和torch.utils.data包來導入數據。

我們今天要解決的問題是訓練一個模型來區分螞蟻和蜜蜂。我們有螞蟻和蜜蜂的訓練圖像各120張。每一類有75張驗證圖片。通常，如果是從零開始訓練，這是一個非常小的數據集。因為我們要使用遷移學習，所以我們的例子應該具有很好地代表性。

這個數據集是一個非常小的圖像子集。

你可以從這里下載數據並解壓到當前目錄。

# 訓練數據的擴充及標准化
# 只進行標准化驗證
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_size = {x:len(image_datasets[x]) for x in ['train','val']}
class_name = image_datasets['train'].classes

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

可視化一些圖像

為了理解數據擴充，我們可視化一些訓練圖像。

def imshow(inp, title=None):
    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(10)    # 暫停一會，以便更新繪圖


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

# 從批處理中生成網格
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_name[x] for x in classes])

訓練模型

現在我們來實現一個通用函數來訓練一個模型。在這個函數中，我們將：

• 調整學習率
• 保存最優模型

下面例子中，參數schedule是來自torch.optim.lr_scheduler的LR調度對象。

def train_model(model, criterion, optimizer, schduler, 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)

        for phase in ['train', 'val']:
            if phase == 'train':
                schduler.step()
                model.train()   # 訓練模型
            else:
                model.eval()    # 評估模型

            running_loss = 0.0
            running_corrects = 0

            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()

                # 統計
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

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

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

            # 很拷貝模型
            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

可視化模型預測

展示部分預測圖像的通用函數：

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_name[preds[j]]))
                imshow(inputs.cpu().data[j])

                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return

        model.train(mode=was_training)

Finetuning the convnet

加載預處理的模型和重置最后的全連接層：

model_ft = models.resnet18(pretrained=True)
num_ftrs = model_ft.fc.in_features
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次，學習率衰減0.1
exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer_ft, step_size=7, gamma=0.1)

訓練和評估

在CPU上可能會花費15-25分鍾，但是在GPU上，少於1分鍾。

model_ft = train_model(model_ft, criterion, optimizer_ft,
                       exp_lr_scheduler, num_epochs=25)

visualize_model(model_ft)

ConvNet作為固定特征提取器

現在，我們凍結除最后一層外的所有網絡。我們需要設置requires_grad=False來凍結參數，這樣調用backward()時不計算梯度。

你可以從這篇文檔中了解更多。

model_conv = 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_ft = optim.SGD(model_conv.parameters(), lr=0.001, momentum=0.9)

# 沒7次，學習率衰減0.1
exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer_ft, step_size=7, gamma=0.1)

model_conv = train_model(model_conv, criterion, optimizer_ft,
                         exp_lr_scheduler, num_epochs=25)

visualize_model(model_conv)

plt.ioff()
plt.show()