Pytorch實現CNN模型的遷移學習——蜜蜂和螞蟻圖片分類項目

很多時候當訓練一個新的圖像分類任務時，一般不會完全從一個隨機的模型開始訓練，而是利用預訓練的模型來加速訓練的過程。經常使用在ImageNet上的預訓練模型。

• 這是一種transfer learning的方法。我們常用以下兩種方法做遷移學習。
  • fine tuning: 從一個預訓練模型開始，我們改變一些模型的架構，然後繼續訓練整個模型的參數。
  • feature extraction: 我們不再改變與訓練模型的參數，而是隻更新我們改變過的部分模型參數。我們之所以叫它feature extraction是因爲我們把預訓練的CNN模型當做一個特徵提取模型，利用提取出來的特徵做來完成我們的訓練任務。
    以下是構建和訓練遷移學習模型的基本步驟：
• 初始化預訓練模型
• 把最後一層的輸出層改變成我們想要分的類別總數
• 定義一個optimizer來更新參數
• 模型訓練

一、項目介紹：

1、目標：本文嘗試採用CNN實現圖像蜜蜂和螞蟻圖像分類任務。
2、數據說明：使用hymenoptera_data數據集。包括兩類圖片, bees 和 ants, 這些數據都被處理成了可以使用ImageFolder來讀取的格式。
輸入數據維度：；輸出數據維度
3、torchvision的datasets.ImageFolder參數說明：
（1）data_dir：數據的存儲目錄。設置成數據的根目錄
（2）model_name：訓練時使用的模型。設置成自己的訓練模型，也可以使用封裝好的模型。如，resnet, alexnet, vgg, squeezenet, densenet, inception等
（3）num_classes：表示數據集分類的類別數
（4）batch_size
（5）num_epochs
（6）feature_extract：表示訓練時使用fine tuning還是feature extraction方法。如果feature_extract = False，整個模型都會被同時更新。如果feature_extract = True，只有模型的最後一層被更新。
4、網絡框架：

二、設置參數

import numpy as np
import torchvision
from torchvision import datasets, transforms, models

import matplotlib.pyplot as plt
import time
import os
import copy
print("Torchvision Version: ",torchvision.__version__)

data_dir = './hymenoptera_data'
#模型可以選擇[resnet, alexnet, vgg, squeezenet, densenet, inception]
model_name = 'resnet'
#數據的label分類個數
num_classes = 2
#訓練的batch_size
batch_size=32
#訓練的輪數
num_epochs = 15
#整個模型參數都參數訓練，進行更新
feature_extract = True

二、代碼實現

1、初始化模型

#定義需要更新的參數
def set_parameter_requires_grad(model,feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False
#使用resnet框架進行模型初始化           
def initialize_model(model_name,num_class,feature_extract,use_pretrained=True):
    if model_name == 'renet':
        model_ft = models.resnet18(pretrained=use_pretrained)  #使用resnet18最爲初始化模型
        set_parameter_requires_grad(model_ft,feature_extract)  #模型中所有參數都更新
        num_ftrs = model_ft.fc.in_features
        model_ft.fc = nn.Linear(num_ftrs,num_classes)  #根據分類個數重新定義最後一層全連接層
        input_size = 224
        
    return model_ft,input_size
model_ft,input_size = initialize_model(model_name,num_classes,feature_extract,use_pretrained=True)
print(model_ft) 

2、導入數據集

根據模型輸入的size,將數據預處理稱爲對應的格式

#定義圖片的處理方式
data_transforms = {
    'train':transforms.Compose([
        transforms.RandomResizedCrop(input_size),
        transforms.RandomHorizontalFlip(),
        tranforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  #對三個通道分別從用不同的形式進行標準化
    ]),
    'val':transforms.Compose([
        transforms.Resize(input_size),
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
        transforms.Normalize([[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]])
    ])
}

print("Initializing Datasets and Dataloaders...")

#生成訓練集和測試集
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir,x),data_transforms[x]) for x in ['train','val']}
#生成訓練集和測試集的迭代器
dataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x],batch_size=batch_size,shuffle=True,num_workers=4) for x in ['train','val']}

#判斷是否使用gpu
device = torch.device('cuda:0' if torch.cuda.is_availabel() else 'cpu')

3、觀察原始數據

it = iter(dataloaders_dict["train"])
inputs, labels = next(it)
for inputs, labels in dataloaders_dict["train"]:
    print(labels.size())
    
print(len(dataloaders_dict["train"].dataset.imgs),len(dataloaders_dict["train"].dataset))

4、模型訓練

torch.manual_seed(53113
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
                  
def train_model(model,dataloaders,criterion,optimizer,num_epochs=5):
    since = time.time()
    val_acc_history = []
    best_model_wts = copy.deepcopy(model.state_dict())   #保存最佳參數；copy.deepcopy一旦複製出來了，就應該是獨立的了，原始變量改變其不會變，數值一樣，但對象不同的。
    best_acc = 0.  #保存最高的準確率
    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs-1}')
        print("-"*10)
        
        for phase in ['train','val']:
            running_loss = 0.
            running_corrects = 0.
            if phase == 'train':
                model.train()
            else:
                model.eval()
                
            for inputs, labels in dataloaders[phase]:   
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                with torch.autograd.set_grad_enabled(phase=='train'):   #當phase=='train'爲True是，可以進行自動求梯度計算。
                    outputs = model(inputs)   #將數據輸入到模型中進行預測
                    loss = criterion(outputs,labels)  #損失計算
                    
                _,preds = torch.max(outputs,1)   #在所有輸出中選取最大的最爲輸出概率
                if phase == 'train':
                    optimizer.zero_grad()   #每一次梯度計算後，所有梯度設置爲0
                    loss.backward()    #反向求導
                    optimizer.step()   #參數優化
                    
                running_loss += loss.item() * inputs.size(0)   #損失*batch，等於每一批的總損失；每一批的總損失相加爲所有樣本的損失和
                running_corrects += torch.sum(preds.view(-1) == labels.view(-1))  #每批的預測正確的樣本數量；每批預測正確的樣本數量相加爲所有樣本中預測正確的樣本數量
            
            epoch_loss = running_loss / len(dataloaders[phase].dataset)  #樣本總損失 / 樣本總個數 = 平均損失
            epoch_acc = running_corrects / len(dataloaders[phase].dataset) #樣本預測正確的個數 / 樣本總個數 = 平均準確率
                
            print(f'{phase} Loss:{epoch_loss} Acc:{epoch_acc}')  #取平均準確率最爲當前epoch的準確率，與之前的進行比較，保存最高的
            if phase == 'val' and epoch_acc > best_acc:
                  best_acc = epoch_acc  #每個epoch更新一次準確率
                  best_model_wts = copy.deepcopy(model.state_dict())
            if phase == 'val':
                  val_acc_history.append(epoch_acc)  #將每個epoch的準確率添加到列表中
 
            
    time_elapsed = time.time() - since
    print(f'Training compete in {time_elapsed // 60}m {time_elapsed % 60}s')   #time_elapsed // 60表示相除後取商。
    print(f'Best val Acc: {best_acc}')
    
    model.load_state_dict(best_model_wts)
    return model,val_acc_history   #返回模型和評估集，每個epoch的最佳acc    

（1）對參數進行訓練的訓練集、測試集結果

model_ft = model_ft.to(device)

params_to_update = model_ft.parameters()  #獲取模型所有的參數進行更新
print('Params to learn:')
if feature_extract:  #如果feature_extract爲True,表示對參數進行更新
    params_to_update = []
    for name,param in model_ft.named_parameters():
        if param.requirs_grad == True:
            params_to_update.append(param)
            print('\t',name)
            
else:
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            print('\t',name)

訓練結果：
Training compete in 0.0m 14.700076580047607s
Best val Acc: 0.954248366013072
（2）不對參數訓練的訓練集和測試集結果

scratch_model,_ = initlialize_model(model_name,num_classes,feature_extract=False,use_pretrained=False)
scratch_model = scratch_model.to(device)
scratch_optimizer = optim.SGD(scratch_model.parameters(), lr=0.001, momentum=0.9)
scratch_criterion = nn.CrossEntropyLoss()
-,scratch_hist = train_model(scratch_model, dataloaders_dict, scratch_criterion, scratch_optimizer, num_epochs=num_epochs)

訓練結果：
Training compete in 0.0m 18.418847799301147s
Best val Acc: 0.7450980392156863
（3）是否對參數訓練，效果對比

plt.title("Validation Accuracy vs. Number of Training Epochs")
plt.xlabel("Training Epochs")
plt.ylabel("Validation Accuracy")
plt.plot(range(1,num_epochs+1),ohist,label="Pretrained")
plt.plot(range(1,num_epochs+1),scratch_hist,label="Scratch")
plt.ylim((0,1.))
plt.xticks(np.arange(1, num_epochs+1, 1.0))
plt.legend()
plt.show()

參考文獻：七月在線Pytorch課程