介紹如何在 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)