介紹如何在 PyTorch 深度學習架構下透過遷移式學習，使用 ResNet 預訓練模型分類螞蟻與蜜蜂圖片。
載入必要模組
首先載入一些必要的模組：
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
載入資料
這裡我們將從 ImageNet 中取出螞蟻（ants）與蜜蜂（bees）的圖片各約 120 張作為訓練資料集，然後拿另外各約 75 張圖片作為驗證資料集，這樣的資料量對於一般的類神經網路來說算是非常少的，適合採用遷移式學習的方式來處理。實際的資料可以從 PyTorch 的網站上下載。
這裡我們將使用 torchvision 與 torch.utils.data 來載入資料。
# 資料轉換函數 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]) # 標準化 ]), }
這裡在進行標準化時所使用的平均值 [0.485, 0.456, 0.406] 與標準差 [0.229, 0.224, 0.225] 是從 ImageNet 整個資料集所算出來的值，是否要採用 ImageNet 的平均值與標準差，取決於自己的資料分布是否與 ImageNet 的分布相似，這裡我們的資料就是從 ImageNet 中取出的，所以直接使用即可。
接著使用 datasets.ImageFolder 以目錄名稱為標註資訊，建立 Dataset，並接著建立 DataLoader：
# 資料路徑 data_dir = 'data/hymenoptera_data' # 建立 Dataset image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']} # 建立 DataLoader 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']} print(dataset_sizes)
{'train': 244, 'val': 153}
取得各類別的名稱：
# 取得各類別的名稱 class_names = image_datasets['train'].classes print(class_names)
['ants', 'bees']
若 CUDA 環境可用，則使用 GPU 計算，否則使用 CPU：
# 若 CUDA 環境可用，則使用 GPU 計算，否則使用 CPU device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(device)
cuda:0
查看資料
自行定義一個可將 Tensor 轉為原始圖片的 tensor2img() 函數，他會將標準化之後的 Tensor 影像轉回原始的分布，可用來檢視圖片：
# 將 Tensor 資料轉為原來的圖片 def tensor2img(inp): 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) return inp
從 DataLoader 取得一個 batch 的訓練資料，並顯示出來：
# 取得一個 batch 的訓練資料 inputs, classes = next(iter(dataloaders['train'])) # 將多張圖片拼成一張 grid 圖 out = torchvision.utils.make_grid(inputs) # 顯示圖片 img = tensor2img(out) plt.imshow(img) plt.title([class_names[x] for x in classes])
訓練模型用函數
以下定義一個用於訓練模型的 train_model() 函數，在這個函數中我們使用傳入的模型（model）、損失函數（criterion）、學習優化器（optimizer）、排程器（scheduler）與 epoch 數（num_epochs）來進行模型的訓練，而訓練模型的過程中，最佳的模型參數可能不是發生在最後的疊代結果上，我們以 best_model_wts 與 best_acc 記錄整個訓練過程中最佳的參數組合與準確率。
# 訓練模型用函數 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']: if phase == 'train': model.train() # 將模型設定為訓練模式 else: model.eval() # 將模型設定為驗證模式 running_loss = 0.0 running_corrects = 0 # 以 DataLoader 載入 batch 資料 for inputs, labels in dataloaders[phase]: # 將資料放置於 GPU 或 CPU inputs = inputs.to(device) labels = labels.to(device) # 重設參數梯度（gradient） optimizer.zero_grad() # 只在訓練模式計算參數梯度 with torch.set_grad_enabled(phase == 'train'): # 正向傳播（forward） outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) if phase == 'train': loss.backward() # 反向傳播（backward） optimizer.step() # 更新參數 # 計算統計值 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) if phase == 'train': # 更新 scheduler scheduler.step() 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)) # 記錄最佳模型 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
微調模型
載入一個 ResNet18 的預訓練模型，將最後一層的輸出數量改為圖片的類別數量，這裡的圖片是螞蟻與蜜蜂兩類，所以輸出數量就改為 2。
# 載入 ResNet18 預訓練模型 model_ft = models.resnet18(pretrained=True) # 取得 ResNet18 最後一層的輸入特徵數量 num_ftrs = model_ft.fc.in_features # 將 ResNet18 的最後一層改為只有兩個輸出線性層 # 更一般化的寫法為 nn.Linear(num_ftrs, len(class_names)) model_ft.fc = nn.Linear(num_ftrs, 2) # 將模型放置於 GPU 或 CPU model_ft = model_ft.to(device)
定義損失函數（loss function）以及學習優化器（optimizer），並搭配排程器（scheduler）讓優化器的 learning rate 可以逐步遞減：
# 使用 cross entropy loss criterion = nn.CrossEntropyLoss() # 學習優化器 optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9) # 每 7 個 epochs 將 learning rate 降為原本的 0.1 倍 exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
接著進行整個模型的微調（fine tuning）：
# 訓練模型 model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)
Epoch 0/24 ---------- train Loss: 0.5848 Acc: 0.7049 val Loss: 0.2998 Acc: 0.9085 Epoch 1/24 ---------- train Loss: 0.5208 Acc: 0.7623 val Loss: 0.3132 Acc: 0.8889 Epoch 2/24 ---------- train Loss: 0.6152 Acc: 0.7664 val Loss: 0.2477 Acc: 0.9085 Epoch 3/24 ---------- train Loss: 0.6715 Acc: 0.7418 val Loss: 0.4634 Acc: 0.8758 Epoch 4/24 ---------- train Loss: 0.8468 Acc: 0.7705 val Loss: 0.9223 Acc: 0.7451 Epoch 5/24 ---------- train Loss: 0.7059 Acc: 0.7664 val Loss: 0.6195 Acc: 0.8039 Epoch 6/24 ---------- train Loss: 0.6604 Acc: 0.7787 val Loss: 0.5225 Acc: 0.8497 Epoch 7/24 ---------- train Loss: 0.3862 Acc: 0.8402 val Loss: 0.3926 Acc: 0.8889 Epoch 8/24 ---------- train Loss: 0.3900 Acc: 0.8279 val Loss: 0.3755 Acc: 0.8562 Epoch 9/24 ---------- train Loss: 0.3492 Acc: 0.8238 val Loss: 0.3567 Acc: 0.8758 Epoch 10/24 ---------- train Loss: 0.3770 Acc: 0.8443 val Loss: 0.3173 Acc: 0.8627 Epoch 11/24 ---------- train Loss: 0.3617 Acc: 0.8361 val Loss: 0.3286 Acc: 0.8693 Epoch 12/24 ---------- train Loss: 0.3593 Acc: 0.8484 val Loss: 0.3716 Acc: 0.8562 Epoch 13/24 ---------- train Loss: 0.2497 Acc: 0.9016 val Loss: 0.2862 Acc: 0.9216 Epoch 14/24 ---------- train Loss: 0.2342 Acc: 0.9057 val Loss: 0.2898 Acc: 0.9020 Epoch 15/24 ---------- train Loss: 0.2945 Acc: 0.8893 val Loss: 0.2890 Acc: 0.8954 Epoch 16/24 ---------- train Loss: 0.3136 Acc: 0.8730 val Loss: 0.2956 Acc: 0.9020 Epoch 17/24 ---------- train Loss: 0.3005 Acc: 0.8730 val Loss: 0.3301 Acc: 0.8824 Epoch 18/24 ---------- train Loss: 0.3857 Acc: 0.8320 val Loss: 0.2763 Acc: 0.9150 Epoch 19/24 ---------- train Loss: 0.3054 Acc: 0.8811 val Loss: 0.3491 Acc: 0.8824 Epoch 20/24 ---------- train Loss: 0.3824 Acc: 0.8320 val Loss: 0.2865 Acc: 0.8954 Epoch 21/24 ---------- train Loss: 0.2979 Acc: 0.9016 val Loss: 0.2849 Acc: 0.9085 Epoch 22/24 ---------- train Loss: 0.2810 Acc: 0.8689 val Loss: 0.2842 Acc: 0.9150 Epoch 23/24 ---------- train Loss: 0.3856 Acc: 0.8279 val Loss: 0.2969 Acc: 0.8889 Epoch 24/24 ---------- train Loss: 0.2740 Acc: 0.8689 val Loss: 0.2963 Acc: 0.8824 Training complete in 1m 7s Best val Acc: 0.921569
經過 25 個 epochs 的訓練，得到最佳的模型準確度為 92.16%。
使用模型預測
以下定義一個使用模型進行預測，並顯示結果的函數：
# 使用模型進行預測，並顯示結果 def visualize_model(model, num_images=6): was_training = model.training # 記錄模型之前的模式 model.eval() # 將模型設定為驗證模式 images_so_far = 0 fig = plt.figure() with torch.no_grad(): # 以 DataLoader 載入 batch 資料 for i, (inputs, labels) in enumerate(dataloaders['val']): # 將資料放置於 GPU 或 CPU 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]])) # 將 Tensor 轉為原始圖片 img = tensor2img(inputs.cpu().data[j]) ax.imshow(img) if images_so_far == num_images: model.train(mode=was_training) # 恢復模型之前的模式 return model.train(mode=was_training) # 恢復模型之前的模式
使用自行定義的 visualize_model() 函數來測試模型：
# 以模型進行預測
visualize_model(model_ft)
微調部分模型
另外一種微調模型的方式是將預訓練模型的參數都鎖定住，只針對新加上去的最後一層參數進行訓練：
# 載入 ResNet18 預訓練模型 model_conv = models.resnet18(pretrained=True) # 鎖定 ResNet18 預訓練模型參數 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) # 將模型放置於 GPU 或 CPU model_conv = model_conv.to(device) # 使用 cross entropy loss criterion = nn.CrossEntropyLoss() # 只針對輸出層進行優化 optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9) # 每 7 個 epochs 將 learning rate 降為原本的 0.1 倍 exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)
同樣以自己定義的 train_model() 函數進行部分模型的微調：
# 訓練模型 model_conv = train_model(model_conv, criterion, optimizer_conv, exp_lr_scheduler, num_epochs=25)
Epoch 0/24 ---------- train Loss: 0.6269 Acc: 0.6352 val Loss: 0.2385 Acc: 0.8889 Epoch 1/24 ---------- train Loss: 0.5501 Acc: 0.7377 val Loss: 0.3131 Acc: 0.8758 Epoch 2/24 ---------- train Loss: 0.7399 Acc: 0.7008 val Loss: 0.2005 Acc: 0.9346 Epoch 3/24 ---------- train Loss: 0.4904 Acc: 0.7705 val Loss: 0.4418 Acc: 0.8431 Epoch 4/24 ---------- train Loss: 0.5982 Acc: 0.7541 val Loss: 0.2174 Acc: 0.9216 Epoch 5/24 ---------- train Loss: 0.3671 Acc: 0.8320 val Loss: 0.2176 Acc: 0.9281 Epoch 6/24 ---------- train Loss: 0.4211 Acc: 0.8402 val Loss: 0.1965 Acc: 0.9346 Epoch 7/24 ---------- train Loss: 0.4075 Acc: 0.8361 val Loss: 0.1967 Acc: 0.9346 Epoch 8/24 ---------- train Loss: 0.3253 Acc: 0.8607 val Loss: 0.3270 Acc: 0.8954 Epoch 9/24 ---------- train Loss: 0.3498 Acc: 0.8566 val Loss: 0.2603 Acc: 0.9085 Epoch 10/24 ---------- train Loss: 0.3565 Acc: 0.8525 val Loss: 0.2281 Acc: 0.9216 Epoch 11/24 ---------- train Loss: 0.3154 Acc: 0.8607 val Loss: 0.1917 Acc: 0.9216 Epoch 12/24 ---------- train Loss: 0.4305 Acc: 0.7951 val Loss: 0.1892 Acc: 0.9412 Epoch 13/24 ---------- train Loss: 0.3430 Acc: 0.8484 val Loss: 0.1960 Acc: 0.9281 Epoch 14/24 ---------- train Loss: 0.3569 Acc: 0.8484 val Loss: 0.2226 Acc: 0.9150 Epoch 15/24 ---------- train Loss: 0.3241 Acc: 0.8484 val Loss: 0.2367 Acc: 0.9346 Epoch 16/24 ---------- train Loss: 0.3814 Acc: 0.8074 val Loss: 0.2011 Acc: 0.9281 Epoch 17/24 ---------- train Loss: 0.3405 Acc: 0.8402 val Loss: 0.2052 Acc: 0.9281 Epoch 18/24 ---------- train Loss: 0.3502 Acc: 0.8320 val Loss: 0.1890 Acc: 0.9281 Epoch 19/24 ---------- train Loss: 0.3813 Acc: 0.8525 val Loss: 0.2027 Acc: 0.9281 Epoch 20/24 ---------- train Loss: 0.3414 Acc: 0.8361 val Loss: 0.1875 Acc: 0.9281 Epoch 21/24 ---------- train Loss: 0.3542 Acc: 0.8443 val Loss: 0.1867 Acc: 0.9477 Epoch 22/24 ---------- train Loss: 0.3195 Acc: 0.8648 val Loss: 0.2480 Acc: 0.9085 Epoch 23/24 ---------- train Loss: 0.3395 Acc: 0.8566 val Loss: 0.2661 Acc: 0.9020 Epoch 24/24 ---------- train Loss: 0.3682 Acc: 0.8730 val Loss: 0.2305 Acc: 0.9216 Training complete in 1m 3s Best val Acc: 0.947712
只對最後一層輸出層參數進行訓練，同樣經過 25 個 epoches 的訓練，在這個例子中反而得到較高的準確度 94.77%。
使用自行定義的 visualize_model() 函數來測試模型：
# 以模型進行預測
visualize_model(model_conv)
