迁移学习

# License: BSD
# Author: Sasank Chilamkurthy

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
from PIL import Image
from tempfile import TemporaryDirectory

cudnn.benchmark = True
plt.ion()   

# 训练数据的增强和标准化处理
# 验证集仅进行标准化处理
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),      # 随机缩放裁剪至224x224
        transforms.RandomHorizontalFlip(),      # 随机水平翻转（数据增强）
        transforms.ToTensor(),                  # 转为张量格式
        transforms.Normalize(                   # 标准化处理
            [0.485, 0.456, 0.406],             # ImageNet均值
            [0.229, 0.224, 0.225]              # ImageNet标准差
        )
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),                # 调整大小至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'  # 数据目录路径
# 创建训练集和验证集的ImageFolder数据集
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,               # 批量大小设为4
        shuffle=True,               # 打乱数据顺序（仅训练集）
        num_workers=4               # 使用4个子进程加载数据
    )
    for x in ['train', 'val']
}
# 计算数据集大小
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
# 获取类别名称
class_names = image_datasets['train'].classes

# 检测并选择可用的硬件加速器（CUDA/MPS/MTIA/XPU）
# 如果当前有可用的加速器则使用，否则使用CPU
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"使用 {device} 设备")  # 打印当前使用的设备

def imshow(inp, title=None):
    """显示张量格式的图片"""
    inp = inp.numpy().transpose((1, 2, 0))  # 调整维度顺序
    inp = np.clip(inp * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406], 0, 1)  # 反标准化
    plt.imshow(inp)
    if title:
        plt.title(title)
    plt.pause(0.001)

# 显示一个batch的训练图片和标签
inputs, classes = next(iter(dataloaders['train']))
imshow(torchvision.utils.make_grid(inputs), [class_names[x] for x in classes])

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    # 临时保存最佳模型
    with TemporaryDirectory() as tempdir:
        best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')
        torch.save(model.state_dict(), best_model_params_path)
        best_acc = 0.0

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

            # 训练/验证双阶段
            for phase in ['train', 'val']:
                if phase == 'train':
                    model.train()  # 训练模式
                else:
                    model.eval()   # 评估模式

                running_loss = 0.0
                running_corrects = 0

                # 批次处理
                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)  # 获取预测结果
                        loss = criterion(outputs, labels)

                        # 反向传播仅训练阶段执行
                        if phase == 'train':
                            loss.backward()
                            optimizer.step()

                    # 累计统计
                    running_loss += loss.item() * inputs.size(0)
                    running_corrects += torch.sum(preds == labels.data)
                
                # 学习率调整
                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
                    torch.save(model.state_dict(), best_model_params_path)

            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}')

        # 加载最佳模型参数
        model.load_state_dict(torch.load(best_model_params_path, weights_only=True))
    return model

def visualize_model(model, num_images=6):
    was_training = model.training  # 保存原始训练状态
    model.eval()  # 切换为评估模式
    
    images_so_far = 0  # 已显示图像计数
    fig = plt.figure()  # 创建图像窗口

    with torch.no_grad():  # 禁用梯度计算
        for i, (inputs, labels) in enumerate(dataloaders['val']):  # 遍历验证集
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)  # 前向传播
            _, preds = torch.max(outputs, 1)  # 获取预测类别

            # 显示当前batch中的图像
            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(f'predicted: {class_names[preds[j]]}')  # 显示预测标签
                imshow(inputs.cpu().data[j])  # 显示图像

                # 达到指定数量后恢复状态并返回
                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return
                    
        model.train(mode=was_training)  # 恢复原始训练状态

# 加载预训练的ResNet18模型(ImageNet权重)
model_ft = models.resnet18(weights='IMAGENET1K_V1')

# 获取原模型全连接层输入特征数
num_ftrs = model_ft.fc.in_features

# 替换最后的全连接层(适配当前二分类任务)
model_ft.fc = nn.Linear(num_ftrs, 2)  # 输出2个类别(蚂蚁/蜜蜂)

# 将模型转移到指定设备(GPU/CPU)
model_ft = model_ft.to(device)

# 定义交叉熵损失函数
criterion = nn.CrossEntropyLoss()

# 初始化SGD优化器(优化所有参数)
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# 定义学习率调度器(每7个epoch衰减0.1倍)
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)

visualize_model(model_ft)

# 加载预训练ResNet18模型(固定所有底层参数)
model_conv = torchvision.models.resnet18(weights='IMAGENET1K_V1')
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()

# 仅优化全连接层参数(与之前不同)
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)

# 学习率每7个epoch衰减为0.1倍
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)

visualize_model(model_conv)

plt.ioff()
plt.show()

def visualize_model_predictions(model, img_path):
    was_training = model.training  # 保存模型当前训练状态
    model.eval()  # 切换到评估模式

    # 加载并预处理输入图像
    img = Image.open(img_path)  # 读取图像文件
    img = data_transforms['val'](img)  # 应用验证集的数据变换
    img = img.unsqueeze(0)  # 增加batch维度 (C,H,W) -> (1,C,H,W)
    img = img.to(device)  # 转移到计算设备

    # 模型预测
    with torch.no_grad():  # 禁用梯度计算
        outputs = model(img)  # 前向传播
        _, preds = torch.max(outputs, 1)  # 获取预测类别

        # 可视化预测结果
        ax = plt.subplot(2,2,1)  # 创建子图
        ax.axis('off')  # 隐藏坐标轴
        ax.set_title(f'Predicted: {class_names[preds[0]]}')  # 显示预测标签
        imshow(img.cpu().data[0])  # 显示图像

    model.train(mode=was_training)  # 恢复原始训练状态


visualize_model_predictions(
    model_conv,
    img_path='data/hymenoptera_data/val/bees/72100438_73de9f17af.jpg'
)

plt.ioff()
plt.show()