轉(zhuǎn)移學習是一種非常強大的技術(shù)，它允許我們利用預(yù)訓練模型的知識來解決新的問題。本教程將教你如何使用 PyTorch 實現(xiàn)轉(zhuǎn)移學習，進行圖像分類任務(wù)。通過在小型數(shù)據(jù)集上微調(diào)預(yù)訓練模型，即使數(shù)據(jù)量有限，也能取得很好的效果。示例代碼中的注釋和變量名都已翻譯為中文，方便國內(nèi)讀者理解。

一、什么是轉(zhuǎn)移學習？

轉(zhuǎn)移學習是一種機器學習方法，它允許模型在解決一個任務(wù)后，將所學的知識遷移到另一個相關(guān)任務(wù)上。在深度學習中，我們通常在大規(guī)模數(shù)據(jù)集（如 ImageNet）上預(yù)訓練一個模型，然后將其用于新的任務(wù)。具體來說，這里有兩種主要的策略：

• 微調(diào)卷積網(wǎng)絡(luò) ：用預(yù)訓練的網(wǎng)絡(luò)初始化模型，例如在 ImageNet 上訓練的模型，然后在新的數(shù)據(jù)集上繼續(xù)訓練整個網(wǎng)絡(luò)。
• 卷積網(wǎng)絡(luò)作為固定特征提取器 ：凍結(jié)網(wǎng)絡(luò)中除最后一層之外的所有層的權(quán)重，替換最后一層為新的分類層，并只訓練這個新的分類層。

二、載入數(shù)據(jù)

我們將使用一個小型的數(shù)據(jù)集，包含螞蟻和蜜蜂的圖像。數(shù)據(jù)集分為訓練集和驗證集。我們使用 torchvision 和 torch.utils.data 包來加載數(shù)據(jù)，并進行數(shù)據(jù)增強和歸一化處理。

import torch
import torchvision
from torchvision import datasets, models, transforms
import os


## 數(shù)據(jù)增強和歸一化處理
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),  # 隨機裁剪到 224x224
        transforms.RandomHorizontalFlip(),  # 隨機水平翻轉(zhuǎn)
        transforms.ToTensor(),  # 轉(zhuǎn)換為張量
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 歸一化
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),  # 調(diào)整大小到 256x256
        transforms.CenterCrop(224),  # 中心裁剪到 224x224
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}


data_dir = 'data/hymenoptera_data'  # 數(shù)據(jù)集路徑
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_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes  # 類別名稱
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 使用 GPU 或 CPU

三、可視化一些圖像

為了更好地理解數(shù)據(jù)增強的效果，我們可視化一些訓練圖像。

import matplotlib.pyplot as plt
import numpy as np


def imshow(inp, title=None):
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))  # 轉(zhuǎn)換張量維度
    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(0.001)  # 暫停一下以便更新圖像顯示


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


## 制作圖像網(wǎng)格并顯示
out = torchvision.utils.make_grid(inputs)
imshow(out, title=[class_names[x] for x in classes])

四、訓練模型

我們編寫一個通用的函數(shù)來訓練模型，包括設(shè)置學習率調(diào)度、保存最佳模型等功能。

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())  # 保存最佳模型權(quán)重
    best_acc = 0.0


    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)


        # 每個 epoch 包括訓練和驗證兩個階段
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # 設(shè)置為訓練模式
            else:
                model.eval()   # 設(shè)置為評估模式


            running_loss = 0.0
            running_corrects = 0


            # 迭代數(shù)據(jù)
            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)  # 獲取預(yù)測類別
                    loss = criterion(outputs, labels)  # 計算損失


                    # 反向傳播和優(yōu)化（僅在訓練階段）
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()


                # 統(tǒng)計損失和準確率
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)


            # 調(diào)整學習率（僅在訓練階段）
            if phase == 'train':
                scheduler.step()


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


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


            # 深拷貝最佳模型
            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(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'Best val Acc: {best_acc:.4f}')


    # 加載最佳模型權(quán)重
    model.load_state_dict(best_model_wts)
    return model

五、微調(diào)卷積網(wǎng)絡(luò)

加載預(yù)訓練的 ResNet-18 模型，并替換最后一層以適應(yīng)我們的分類任務(wù)。

model_ft = models.resnet18(pretrained=True)  # 加載預(yù)訓練模型
num_ftrs = model_ft.fc.in_features  # 獲取全連接層的輸入特征數(shù)
model_ft.fc = nn.Linear(num_ftrs, 2)  # 替換全連接層為新的分類層（這里假設(shè)有兩類）
model_ft = model_ft.to(device)  # 將模型移動到設(shè)備上


criterion = nn.CrossEntropyLoss()  # 定義損失函數(shù)


## 優(yōu)化所有參數(shù)
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)


## 學習率調(diào)度
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)


## 訓練模型
model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)

六、將卷積網(wǎng)絡(luò)作為固定特征提取器

凍結(jié)預(yù)訓練模型的大部分層，只訓練新的分類層。

model_conv = torchvision.models.resnet18(pretrained=True)  # 加載預(yù)訓練模型


## 凍結(jié)所有參數(shù)
for param in model_conv.parameters():
    param.requires_grad = False


## 替換全連接層
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)


model_conv = model_conv.to(device)


criterion = nn.CrossEntropyLoss()


## 僅優(yōu)化分類層的參數(shù)
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)


## 學習率調(diào)度
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)


## 訓練模型
model_conv = train_model(model_conv, criterion, optimizer_conv, exp_lr_scheduler, num_epochs=25)

七、總結(jié)

通過本教程，你學會了如何在 PyTorch 中利用轉(zhuǎn)移學習進行圖像分類。無論是微調(diào)預(yù)訓練模型還是將其作為固定特征提取器，都能有效地提高模型在小型數(shù)據(jù)集上的性能。