PyTorch練手項目三：模型微調

本文目的：基於kaggle上狗的種類識別項目，展示如何利用PyTorch來進行模型微調。

PyTorch中torchvision是一個針對視覺領域的工具庫，除了提供有大量的數據集，還有許多預訓練的經典模型。這里以官方訓練好的resnet50為例，拿來參加kaggle上面的dog breed狗的種類識別。

1 導入相關庫，設置一些超參

import torch
import torchvision
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, models, transforms
import pandas as pd
import os
from PIL import Image
from sklearn.model_selection import StratifiedShuffleSplit

print(torch.__version__)  #1.1.0
print(torchvision.__version__) #0.3.0


#定義一些超參
IMG_SIZE = 224 #模型要求的輸入尺寸
IMG_MEAN = [0.485, 0.456, 0.406] #圖像預處理中需要的均值和方差
IMG_STD = [0.229, 0.224, 0.225]
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #盡量使用GPU
BATCH_SIZE = 64  #每一個batch的大小
EPOCHS = 7  #訓練輪數

2 准備數據

Pytorch中數據的讀取通常需要封裝成Dataset類對象和DataLoader類對象。

2.1 獲取數據並整理

首先下載官方的數據並解壓，只要保持數據的目錄結構即可，這里指定一下目錄的位置，並且看下內容。（注意：labels.csv文件中有10222條標簽，對應的是train文件夾中圖像。）

#DATA_ROOT = r'D:\KaggleDatasets\competitions\dog-breed-identification'
#注1：常用'/'表相對路徑，'\'表絕對路徑，網頁網址和linux系統下一般用'/'
DATA_ROOT = '/KaggleDatasets/competitions/dog-breed-identification'
df = pd.read_csv(os.path.join(DATA_ROOT, 'labels.csv'))
df.head()

為了后續方便，這里定義兩個字典，並將類別序號添加進DataFrame中。

#分別以標簽字符串和序號為索引，定義兩個字典
breeds = df.breed.unique()
breed2idx = dict((breed,idx) for idx,breed in enumerate(breeds)) 
idx2breed = dict((idx,breed) for idx,breed in enumerate(breeds))
len(breeds) #120

#將類別序號添加到df的列 
df['label_idx'] = pd.Series(breed2idx, index=df.breed).values  
#df.shape  #(10222, 3)
df.head()

將數據分割成訓練集和驗證集。這里只分割10%的數據作為訓練時的驗證數據。

#分割數據集
shuffle_split = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=0) #分層切割
train_idx, val_idx = next(iter(shuffle_split.split(df, df.breed))) #split方法返回迭代器
train_df = df.iloc[train_idx].reset_index(drop=True) #(9199, 3)
val_df = df.iloc[val_idx].reset_index(drop=True)  #(1023, 3)

注2：StratifiedShuffleSplit().split(X, y)

• 含義：This cross-validation object is a merge of StratifiedKFold and ShuffleSplit
• split(X, y)方法中，X實際上只用了 np.zeros(n_samples) 來產生占位符，而切分主要靠y
• KFold(shuffle=True)與ShuffleSplit的區別是：前者只shuffle一次，后者每次切分之前都要shuffle

注3：sklearn中幾種數據切分方法

• train_test_split：普通切分
• KFold：普通K折切分
• StratifiedKFold：分層K折切分
• StratifiedShuffleSplit：每次shuffle后分層切分

2.2 自定義Dataset

torch.utils.data.Dataset是一個抽象類, 自定義的Dataset需要繼承它並且實現兩個成員方法：

• __ len __ () ：返回整個數據集的長度。
• __ getitem __ () ：每次怎么讀取數據。

另外，transform過程也在此處傳進來。

#自定義Dataset
class DogDataset(Dataset):
    def __init__(self, df, img_path, transform=None):
        self.df = df
        self.img_path = img_path
        self.transform = transform
    
    def __len__(self):
        return self.df.shape[0]  #返回數據集長度
    
    def __getitem__(self, idx):  #每次根據idx返回一個(image,label)數據對
        img_name = os.path.join(self.img_path, self.df.id[idx]) + '.jpg'
        img = Image.open(img_name)  #建議用PIL，而非skimage
        label = self.df.label_idx[idx]
        
        if self.transform:
            img = self.transform(img)
        return img, label
    

#自定義訓練集和驗證集的transform
train_transform = transforms.Compose([
    transforms.Resize(IMG_SIZE),
    transforms.RandomResizedCrop(IMG_SIZE),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(30),
    transforms.ToTensor(),
    transforms.Normalize(IMG_MEAN, IMG_STD),
])

test_transform = transforms.Compose([
    transforms.Resize(IMG_SIZE), #注4：傳入一個int時，短邊縮放到IMG_SIZE，長邊按比例縮放
    transforms.CenterCrop(IMG_SIZE),  
    transforms.ToTensor(),
    transforms.Normalize(IMG_MEAN, IMG_STD),
])


#生成dataset
train_dataset = DogDataset(train_df, os.path.join(DATA_ROOT,'train'), train_transform)
val_dataset = DogDataset(val_df, os.path.join(DATA_ROOT,'train'), test_transform)

2.3 定義DataLoader

類定義為：

torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, ...)

可以看到主要參數有這么幾個:

• dataset：即上面自定義的dataset；
• batch_size：一個batch中樣本個數；
• shuffle：划分batch前是否打亂順序；
• sampler：定義抽樣的策略；
• batch_sampler：定義批次抽樣的策略；
• num_worker：定義多線程方法，默認為0。

#生成dataloader
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=True)

3 准備模型

使用Pytorch中torchvision.models.resnet50。由於ImageNet是識別1000個物體，這里狗的分類一共只有120，所以需要對模型的最后一層全連接層進行微調，將輸出從1000改為120。

#准備模型
model = models.resnet50(pretrained=True) #可用dir(model)查看屬性及方法

#將所有參數凍結
for param in model.parameters(): 
    param.requires_grad = False
print(model.fc)

#修改fc層。可用model.named_parameters()迭代查看具體名稱和參數
num_feature = model.fc.in_features  #獲取fc層的輸入個數
model.fc = nn.Linear(num_feature, len(breeds))  #重新定義fc層
print(model.fc)
#print(model)

#將model移至GPU
model.to(DEVICE)

注5：關於預訓練模型的使用，需要

• 傳入pretrained=True，可加載預訓練權重；
• 模型使用前需要調用model.train()，或者model.eval()來開啟或關閉BN和Dropout等；
• 傳給預訓練模型的圖像應符合：（可見2.2中定義的transform）
  • 3通道RGB格式；
  • shape為(3,H,W)，其中H和W至少為224，若不夠則需要Resize；
  • 以[0,1]范圍加載后用mean=[0.485,0.456,0.406]和std=[0.229, 0.224, 0.225]來Normalize

4 訓練

4.1 定義訓練參數和函數

訓練需要定義損失函數和優化器。另外也打包定義了訓練和驗證函數。

#指定損失函數和優化器
loss_fn = nn.CrossEntropyLoss()  #注6：默認的reduction為mean，即求平均損失
#optimizer = torch.optim.Adam([{'params':model.fc.parameters()}], lr=0.001) #定義fc層學習率
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)


#定義訓練函數
#注7：訓練5部曲：梯度清零，前向傳播，計算損失，反向傳播，梯度更新。
def train(model, train_loader, device, epoch):
    model.train()  #注8：開啟訓練模型，即開啟BN和Dropout等
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device) #注9：模型和數據均要移至GPU
        #data和target的size分別為torch.Size([64, 3, 224, 224])、torch.Size([64])
        optimizer.zero_grad() #梯度清零
        yhat = model(data) #前向傳播 torch.Size([64, 120])
        loss = loss_fn(yhat, target) #計算損失
        loss.backward() #反向傳播
        optimizer.step() #更新梯度
    print('Train epoch {}\t Loss {:.6f}'.format(epoch, loss.item()))
    
    
#定義測試函數
def test(model, val_loader, device):
    model.eval()
    test_loss = 0  #記錄測試損失
    correct = 0  #記錄預測正確個數
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(val_loader):
            data, target = data.to(device), target.to(device)
            yhat = model(data)
            test_loss += loss_fn(yhat, target).item() #每次加上一個batch的平均損失值
            pred = torch.max(yhat, dim=1, keepdim=True)[1]  #注10：找到概率最大的下標
            correct += pred.eq(target.view_as(pred)).sum().item() #累加正確的樣本個數
            
    test_loss /= len(val_loader) #注意此處是除以batch個數，而非len(val_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.1f}%)\n'.format(
        test_loss, correct, len(val_loader.dataset),
        100. * correct / len(val_loader.dataset)))

4.2 開始訓練

#開始訓練
for epoch in range(1, EPOCHS+1):
    %time train(model, train_loader, DEVICE, epoch)
    test(model, val_loader, DEVICE)

從結果可以看出，運行幾輪之后准確率大約在80%左右，比隨機猜測（0.83%）要好很多。

Train epoch 1	 Loss 1.935438
Wall time: 3min 26s

Test set: Average loss: 1.2672, Accuracy: 723/1023 (70.7%)

Train epoch 2	 Loss 1.673698
Wall time: 1min 41s

Test set: Average loss: 0.8607, Accuracy: 782/1023 (76.4%)

Train epoch 3	 Loss 1.657430
Wall time: 1min 41s

Test set: Average loss: 0.7643, Accuracy: 795/1023 (77.7%)

Train epoch 4	 Loss 1.463368
Wall time: 1min 40s

Test set: Average loss: 0.7109, Accuracy: 806/1023 (78.8%)

Train epoch 5	 Loss 1.849077
Wall time: 1min 40s

Test set: Average loss: 0.7227, Accuracy: 803/1023 (78.5%)

Train epoch 6	 Loss 1.442590
Wall time: 1min 40s

Test set: Average loss: 0.7080, Accuracy: 796/1023 (77.8%)

Train epoch 7	 Loss 1.540823
Wall time: 1min 41s

Test set: Average loss: 0.6738, Accuracy: 822/1023 (80.4%)

5 小結

• 普通任務的過程：准備數據、准備模型、訓練、評估或預測；
• 如何對預訓練模型進行微調；
• 利用Pandas和sklearn工具處理數據；
• 標注的10個注意事項。

Reference

• https://github.com/zergtant/pytorch-handbook
• https://pytorch.org/docs/stable/index.html