PyTorch实战:基于 PyTorch 和 ResNet 预训练模型的迁移学习实战(猫狗分类)

Langchain系列文章目录

01-玩转LangChain:从模型调用到Prompt模板与输出解析的完整指南
02-玩转 LangChain Memory 模块:四种记忆类型详解及应用场景全覆盖
03-全面掌握 LangChain:从核心链条构建到动态任务分配的实战指南
04-玩转 LangChain:从文档加载到高效问答系统构建的全程实战
05-玩转 LangChain:深度评估问答系统的三种高效方法(示例生成、手动评估与LLM辅助评估)
06-从 0 到 1 掌握 LangChain Agents:自定义工具 + LLM 打造智能工作流!

系列文章目录

Pytorch基础篇

01-PyTorch新手必看:张量是什么?5 分钟教你快速创建张量!
02-张量运算真简单!PyTorch 数值计算操作完全指南
03-Numpy 还是 PyTorch?张量与 Numpy 的神奇转换技巧
04-揭秘数据处理神器:PyTorch 张量拼接与拆分实用技巧
05-深度学习从索引开始:PyTorch 张量索引与切片最全解析
06-张量形状任意改!PyTorch reshape、transpose 操作超详细教程
07-深入解读 PyTorch 张量运算:6 大核心函数全面解析,代码示例一步到位!
08-自动微分到底有多强?PyTorch 自动求导机制深度解析

Pytorch实战篇

09-从零手写线性回归模型:PyTorch 实现深度学习入门教程
10-PyTorch 框架实现线性回归:从数据预处理到模型训练全流程
11-PyTorch 框架实现逻辑回归:从数据预处理到模型训练全流程
12-PyTorch 框架实现多层感知机(MLP):手写数字分类全流程详解
13-PyTorch 时间序列与信号处理全解析:从预测到生成
14-深度学习必备:PyTorch数据加载与预处理全解析
15-PyTorch实战:手把手教你完成MNIST手写数字识别任务
16-PyTorch 训练循环全攻略:从零到精通的深度学习秘籍
17-PyTorch实现CNN:CIFAR-10图像分类实战教程
18-RNN 实战指南:用 PyTorch 从零实现文本分类
19-PyTorch实战:基于 PyTorch 和 ResNet 预训练模型的迁移学习实战(猫狗分类)


文章目录

  • Langchain系列文章目录
  • 系列文章目录
  • 前言
  • 一、 什么是迁移学习
    • 1.1 核心概念
    • 1.2 为何需要迁移学习
    • 1.3 迁移学习的主要方式
  • 二、 预训练模型:巨人的肩膀
    • 2.1 什么是预训练模型
    • 2.2 ResNet 模型简介
    • 2.3 为何选择 ResNet
  • 三、 实战:使用 ResNet 微调进行猫狗分类
    • 3.1 环境准备与数据加载
      • 3.1.1 安装 PyTorch 和 TorchVision
      • 3.1.2 数据集准备
      • 3.1.3 数据预处理与增强
      • 3.1.4 创建 DataLoader
    • 3.2 加载预训练 ResNet 模型
      • 3.2.1 导入模型
      • 3.2.2 理解模型结构
    • 3.3 修改模型以适应新任务
      • 3.3.1 冻结部分层参数 (可选但常用)
      • 3.3.2 替换或修改分类层
    • 3.4 定义损失函数和优化器
      • 3.4.1 损失函数
      • 3.4.2 优化器
    • 3.5 训练与验证模型
      • 3.5.1 训练循环
      • 3.5.2 启动训练
      • 3.5.3 模型保存
    • 3.6 完整代码示例 (整合)
  • 四、 常见问题与进阶技巧
    • 4.1 如何选择合适的预训练模型?
    • 4.2 如何设置学习率?
    • 4.3 过拟合与欠拟合怎么办?
    • 4.4 进一步提升性能?
  • 五、 总结


前言

在深度学习领域,尤其是在计算机视觉任务中,从零开始训练一个高性能的模型往往需要海量的数据和强大的计算资源。然而,现实项目中我们常常面临数据量有限、标注成本高昂或训练时间紧迫等挑战。这时,“迁移学习”(Transfer Learning)便成为了一个强大而高效的解决方案。它允许我们站在“巨人”的肩膀上,利用在大规模数据集(如 ImageNet)上预训练好的模型作为起点,针对我们自己的特定任务进行微调(Fine-tuning),从而在数据相对较少的情况下,也能快速构建出性能优异的模型。

本文作为 PyTorch 实战篇 系列的第三篇,将聚焦于迁移学习的核心思想与实践。我们将以经典的“猫狗分类”任务为例,手把手教你如何利用 PyTorch 加载强大的预训练模型 ResNet,并通过微调技术,高效地训练一个准确的图像分类器。无论你是深度学习初学者,还是希望提升模型性能的进阶者,都能从中获益。

一、 什么是迁移学习

1.1 核心概念

迁移学习,顾名思义,就是将从一个任务(源任务)中学到的知识“迁移”到另一个相关但不完全相同的任务(目标任务)上。

想象一下,你学会了骑自行车(源任务),这个过程中你掌握了平衡、转向等基本技能。当你再去学骑摩托车(目标任务)时,虽然摩托车更复杂,但你之前学到的平衡感和转向技巧依然适用,让你能更快上手。这就是迁移学习的直观类比。

在深度学习中,源任务通常是在一个非常庞大的通用数据集(如包含1000类物体的 ImageNet)上训练模型。这个过程中,模型学会了识别图像的各种底层和高层特征,比如边缘、纹理、形状,甚至是一些物体的部件。这些学到的特征对于许多其他的视觉任务(目标任务,如我们的猫狗分类)同样具有很强的泛化能力。

1.2 为何需要迁移学习

迁移学习之所以备受青睐,主要得益于以下几个显著优势:

  1. 数据需求减少: 目标任务不再需要庞大的标注数据集。利用预训练模型学到的通用特征,我们只需要相对少量的数据就能进行有效的微调。
  2. 训练时间缩短: 基于预训练模型的权重进行微调,模型的起点更好,收敛速度通常比从零开始训练快得多。
  3. 模型性能提升: 预训练模型在大数据集上学到的特征表示通常非常强大和鲁棒,有助于提升模型在目标任务上的泛化能力和最终性能,尤其是在目标数据有限时。
  4. 降低计算成本: 缩短了训练时间,也就意味着减少了 GPU/TPU 等计算资源的使用。

1.3 迁移学习的主要方式

迁移学习在实践中主要有两种常见方式:

  1. 特征提取(Feature Extraction): 将预训练模型(通常是卷积层部分)当作一个固定的特征提取器。输入图像通过预训练模型,得到其特征表示(通常是某个中间层或最后卷积层的输出),然后只训练一个新的、简单的分类器(如全连接层)来处理这些提取到的特征,完成目标任务。预训练模型的权重在此过程中保持不变(冻结)。
  2. 微调(Fine-tuning): 加载预训练模型的权重,并替换掉模型原有的、针对源任务的分类头(如最后的全连接层)。然后,使用目标任务的数据对模型的 部分全部 层的权重进行“微调”,使其更适应新任务。通常,我们会“冻结”模型早期的卷积层(它们学习的是通用特征),只微调后面的层和新加的分类层,或者以较小的学习率微调所有层。

本文将重点介绍和实践第二种方式:微调。

二、 预训练模型:巨人的肩膀

2.1 什么是预训练模型

预训练模型(Pre-trained Model)是指已经在一个大型基准数据集上(最著名的是 ImageNet 数据集,包含超过百万张图像和1000个类别)训练好的神经网络模型。这些模型通常由顶尖的研究机构或公司开发和训练,耗费了大量的计算资源,其学到的权重参数蕴含了丰富的通用视觉知识。

常见的图像领域预训练模型包括:

  • AlexNet
  • VGG (VGG16, VGG19)
  • GoogLeNet / Inception
  • ResNet (ResNet18, ResNet34, ResNet50, ResNet101, ResNet152)
  • DenseNet
  • EfficientNet
  • Vision Transformer (ViT)
  • … 等等

PyTorch 通过 torchvision.models 模块提供了许多常用的预训练模型及其权重,方便我们直接加载使用。

2.2 ResNet 模型简介

ResNet(Residual Network,残差网络)是深度学习发展史上的一个里程碑式模型,由 Kaiming He 等人在 2015 年提出。它巧妙地引入了“残差块”(Residual Block)结构,解决了深度神经网络训练中常见的梯度消失/梯度爆炸问题,使得训练非常深的网络(甚至超过1000层)成为可能。

其核心思想是:与其让网络层直接学习目标映射 H ( x ) H(x) H(x),不如让它学习残差映射 F ( x ) = H ( x ) − x F(x) = H(x) - x F(x)=H(x)x,其中 x x x 是该层的输入。原始的目标映射则变为 H ( x ) = F ( x ) + x H(x) = F(x) + x H(x)=F(x)+x。这个 + x + x +x 的操作通过一个“快捷连接”(Shortcut Connection)或“跳跃连接”(Skip Connection)实现,直接将输入 x x x 添加到后面层的输出上。

残差块
快捷连接
卷积层/激活层...
卷积层/激活层...
输入 x
+
输出 H(x)

这种结构使得信息可以直接跨层传播,梯度也更容易回传,极大地提升了深层网络的训练效率和性能。

2.3 为何选择 ResNet

在众多的预训练模型中,ResNet 因其以下优点而成为迁移学习的热门选择:

  • 性能优异: 在 ImageNet 等多个基准数据集上表现出色。
  • 结构灵活: 有不同深度(层数)的版本(如 ResNet18, ResNet34, ResNet50 等),可以根据任务复杂度和计算资源进行选择。层数较少的版本(如 ResNet18/34)计算量相对较小,适合快速实验和资源受限的场景。
  • 广泛可用: 在主流深度学习框架(如 PyTorch, TensorFlow)中都有现成的实现和预训练权重。
  • 效果稳定: 相较于一些更复杂的模型,ResNet 的训练通常更加稳定。

因此,在本实战中,我们选择 ResNet(具体选择如 ResNet18 或 ResNet34,因其效率较高)作为我们的预训练模型。

三、 实战:使用 ResNet 微调进行猫狗分类

接下来,我们将一步步展示如何使用 PyTorch 加载预训练的 ResNet 模型,并对其进行微调,以完成猫狗图像的二分类任务。

3.1 环境准备与数据加载

3.1.1 安装 PyTorch 和 TorchVision

确保你的 Python 环境中安装了 PyTorch 和 TorchVision。如果尚未安装,可以通过 pip 或 conda 安装:

# 使用 pip 安装
pip install torch torchvision torchaudio

# 或者使用 conda 安装 (根据你的 CUDA 版本选择合适的命令,参考 PyTorch 官网)
# conda install pytorch torchvision torchaudio cudatoolkit=x.x -c pytorch

3.1.2 数据集准备

你需要一个包含猫和狗图像的数据集。一个常见的数据集是 Kaggle 上的 “Dogs vs. Cats” 竞赛数据集。你需要将其整理成 PyTorch ImageFolder 能够识别的格式,通常是这样:

/path/to/your/dataset/
├── train/
│   ├── cat/
│   │   ├── cat.0.jpg
│   │   ├── cat.1.jpg
│   │   └── ...
│   └── dog/
│       ├── dog.0.jpg
│       ├── dog.1.jpg
│       └── ...
└── val/  (或者 test)
    ├── cat/
    │   ├── cat.10000.jpg
    │   └── ...
    └── dog/
        ├── dog.10000.jpg
        └── ...

其中,train 目录包含训练图片,val 目录包含验证图片(用于模型选择和评估)。每个目录下再按类别(cat, dog)分子目录存放对应的图片。

3.1.3 数据预处理与增强

使用 torchvision.transforms 来定义数据预处理和数据增强操作。对于预训练模型,关键一步 是使用其在 ImageNet 上训练时所用的相同的归一化参数。

import torch
import torchvision
from torchvision import transforms, datasets
from torch.utils.data import DataLoader

# 定义数据变换
# 注意:预训练模型通常在 ImageNet 数据集上训练,需要使用其均值和标准差进行归一化
# ImageNet 均值和标准差
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224), # 随机裁剪并缩放到 224x224
        transforms.RandomHorizontalFlip(), # 随机水平翻转
        transforms.ToTensor(),             # 转换为 Tensor
        transforms.Normalize(mean, std)    # 归一化
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),            # 缩放到 256x256
        transforms.CenterCrop(224),        # 中心裁剪到 224x224
        transforms.ToTensor(),             # 转换为 Tensor
        transforms.Normalize(mean, std)    # 归一化
    ]),
}

# 数据集路径 (请替换成你的实际路径)
data_dir = '/path/to/your/dataset'
  • RandomResizedCrop(224): 对于训练集,随机裁剪输入图像到 224x224 大小,这是一种常用的数据增强手段。ResNet 通常使用 224x224 的输入。
  • RandomHorizontalFlip(): 随机以 50% 的概率水平翻转图像,增加数据多样性。
  • Resize(256)CenterCrop(224): 对于验证集,通常先将图像等比例缩放到稍大尺寸(如 256x256),然后从中心裁剪出 224x224,以获得稳定的评估结果。
  • ToTensor(): 将 PIL 图像或 NumPy ndarray 转换为 PyTorch Tensor,并将像素值从 [0, 255] 缩放到 [0.0, 1.0]。
  • Normalize(mean, std): 使用 ImageNet 的均值和标准差对图像进行归一化。这是使用预训练模型时非常重要的一步!

3.1.4 创建 DataLoader

使用 datasets.ImageFolder 加载数据,并用 DataLoader 创建数据加载器,以便在训练时按批次(batch)加载数据。

# 加载数据集
image_datasets = {x: datasets.ImageFolder(f"{data_dir}/{x}", data_transforms[x])
                  for x in ['train', 'val']}

# 创建数据加载器
dataloaders = {x: DataLoader(image_datasets[x], batch_size=32, # 可调整 batch size
                             shuffle=True if x == 'train' else False, # 训练集打乱,验证集不打乱
                             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
num_classes = len(class_names)

print(f"训练集大小: {dataset_sizes['train']}")
print(f"验证集大小: {dataset_sizes['val']}")
print(f"类别名称: {class_names}") # 输出应为 ['cat', 'dog'] 或类似
print(f"类别数量: {num_classes}")   # 输出应为 2

3.2 加载预训练 ResNet 模型

3.2.1 导入模型

使用 torchvision.models 来加载预训练的 ResNet 模型。这里我们以 ResNet18 为例。设置 pretrained=True 会自动下载并加载 ImageNet 上的预训练权重。

import torchvision.models as models
import torch.nn as nn

# 加载预训练的 ResNet18 模型
model_ft = models.resnet18(pretrained=True)
# 如果你想用 ResNet34 或 ResNet50,只需替换:
# model_ft = models.resnet34(pretrained=True)
# model_ft = models.resnet50(pretrained=True)

print("原始 ResNet18 模型结构:")
# print(model_ft) # 可以取消注释这行来查看完整结构

3.2.2 理解模型结构

打印模型结构(如取消上面代码的注释)会显示 ResNet 的所有层。你会注意到最后通常有一个名为 fc(全连接)的层,它的输出维度是 1000,对应 ImageNet 的 1000 个类别。我们需要修改这一层以适应我们的猫狗二分类任务。

3.3 修改模型以适应新任务

3.3.1 冻结部分层参数 (可选但常用)

为了保留预训练模型学到的通用特征,并加速训练,我们通常会“冻结”模型早期的卷积层,使其在训练过程中权重不被更新。只训练后面的层和我们新添加的分类层。

# 冻结所有卷积层的参数 (可选策略)
# for param in model_ft.parameters():
#     param.requires_grad = False

# 更精细的策略:只冻结部分层,或完全不冻结(让所有层都微调,但可能需要更小的学习率)
# 例如,冻结除了最后几个 block 之外的所有层
# ct = 0
# for child in model_ft.children():
#     ct += 1
#     if ct < 7: # 假设我们想冻结前 6 个 'child' 模块 (这需要根据具体模型结构调整)
#         for param in child.parameters():
#             param.requires_grad = False

# 简单的策略:先冻结所有,再解冻最后的全连接层(将在下一步替换)
for param in model_ft.parameters():
    param.requires_grad = False

注意: 冻结哪些层是一个可以调整的超参数。如果你的目标任务与 ImageNet 非常相似且数据量充足,可以考虑微调更多层甚至所有层。如果数据量很少,冻结大部分层通常效果更好。

3.3.2 替换或修改分类层

ResNet 原始的 fc 层是为 ImageNet 的 1000 类设计的。我们需要将其替换为一个新的全连接层,其输出维度等于我们的目标任务类别数(猫狗分类为 2)。

# 获取原始全连接层的输入特征数
num_ftrs = model_ft.fc.in_features
print(f"原始 FC 层输入特征数: {num_ftrs}")

# 创建一个新的全连接层,替换掉原来的 fc 层
# 输出维度为我们的类别数 (num_classes = 2)
model_ft.fc = nn.Linear(num_ftrs, num_classes)

print("\n修改后的模型结构 (只看最后部分):")
print(model_ft.fc)

# 确保新的全连接层的参数是可训练的 (如果之前冻结了所有层)
for param in model_ft.fc.parameters():
    param.requires_grad = True

# 将模型移动到 GPU (如果可用)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_ft = model_ft.to(device)
print(f"\n模型已移动到: {device}")

现在,model_ft 就是我们准备好用于微调的模型了。只有最后的全连接层(以及可能解冻的其他层)的参数会在训练中更新。

3.4 定义损失函数和优化器

3.4.1 损失函数

对于多分类(包括二分类)任务,交叉熵损失(Cross Entropy Loss)是标准的损失函数。

criterion = nn.CrossEntropyLoss()

3.4.2 优化器

选择一个优化器,如 SGD 或 Adam。关键点:优化器应该只更新那些 requires_grad=True 的参数。我们可以通过过滤模型的参数列表来实现这一点。同时,为微调设置一个合适的学习率通常比从零训练要小。

import torch.optim as optim
from torch.optim import lr_scheduler

# 定义优化器 - 只优化需要更新的参数
# 将 model_ft.parameters() 替换为只包含 requires_grad=True 的参数列表
params_to_update = model_ft.parameters()
print("需要更新的参数:")
params_to_update_list = []
for name,param in model_ft.named_parameters():
    if param.requires_grad == True:
        params_to_update_list.append(param)
        # print("\t",name) # 取消注释可查看具体哪些参数会被更新

# 使用 Adam 优化器
optimizer_ft = optim.Adam(params_to_update_list, lr=0.001) # 学习率可以调整

# (可选) 添加学习率调度器,例如每隔 N 个 epoch 降低学习率
# exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

3.5 训练与验证模型

现在我们可以编写标准的 PyTorch 训练和验证循环。

3.5.1 训练循环

import time
import copy

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(f'Epoch {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

            # 迭代数据
            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' and scheduler is not None:
            #      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(f'找到更好的模型,验证集准确率: {best_acc:.4f}')


        print()

    time_elapsed = time.time() - since
    print(f'训练完成,耗时 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'最佳验证集准确率: {best_acc:4f}')

    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model

3.5.2 启动训练

调用 train_model 函数开始训练。

# 开始训练模型
model_ft = train_model(model_ft, criterion, optimizer_ft, # exp_lr_scheduler, # 如果使用调度器
                       num_epochs=15) # 可以调整训练轮数

3.5.3 模型保存

训练完成后,model_ft 变量中包含了在验证集上表现最好的模型权重。你可以将其保存到文件,以便后续使用或部署。

# 保存训练好的模型
model_save_path = 'resnet18_finetuned_catdog.pth'
torch.save(model_ft.state_dict(), model_save_path)
print(f"最佳模型已保存至: {model_save_path}")

# 如何加载已保存的模型 (示例)
# model_loaded = models.resnet18(pretrained=False) # 注意这里 pretrained=False
# num_ftrs = model_loaded.fc.in_features
# model_loaded.fc = nn.Linear(num_ftrs, num_classes)
# model_loaded.load_state_dict(torch.load(model_save_path))
# model_loaded = model_loaded.to(device)
# model_loaded.eval() # 设置为评估模式

3.6 完整代码示例 (整合)

# -----------------------------------
# --- 导入必要的库 ---
# -----------------------------------
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torchvision
from torchvision import datasets, models, transforms
import time
import os
import copy

print("PyTorch Version: ", torch.__version__)
print("Torchvision Version: ", torchvision.__version__)

# -----------------------------------
# --- 1. 数据准备与加载 ---
# -----------------------------------
# 数据集路径 (!!!请务必修改为你的实际路径!!!)
data_dir = '/path/to/your/cat_dog_dataset' # 例如 './data/cats_vs_dogs_small'

# 数据变换 (使用 ImageNet 均值和标准差)
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ]),
}

# 创建数据集
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=32,
                                             shuffle=True if x == 'train' else False, 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
num_classes = len(class_names)

# 检查设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"将在 {device} 上进行训练...")
print(f"类别: {class_names}")

# -----------------------------------
# --- 2. 加载和修改预训练模型 ---
# -----------------------------------
# 加载预训练的 ResNet18
model_ft = models.resnet18(pretrained=True)

# (可选) 冻结模型的部分或全部层
freeze_layers = True # 设置为 False 则微调所有层
if freeze_layers:
    for param in model_ft.parameters():
        param.requires_grad = False

# 获取最后一层输入特征数,并替换为新的适应我们任务的层
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)

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

# 打印需要更新的参数 (确认哪些层被冻结/解冻)
print("\n需要更新梯度的参数:")
params_to_update = []
for name, param in model_ft.named_parameters():
    if param.requires_grad:
        params_to_update.append(param)
        print("\t", name)

# -----------------------------------
# --- 3. 定义损失函数和优化器 ---
# -----------------------------------
criterion = nn.CrossEntropyLoss()

# 只优化 requires_grad=True 的参数
optimizer_ft = optim.Adam(params_to_update, lr=0.001) # 学习率可调

# (可选) 学习率调度器
# exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

# -----------------------------------
# --- 4. 训练和验证函数 ---
# -----------------------------------
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
    val_acc_history = [] # 记录验证集准确率历史

    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)
                    loss = criterion(outputs, labels)
                    _, preds = torch.max(outputs, 1)

                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            # # 更新学习率
            # if phase == 'train' and scheduler is not None:
            #     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':
                val_acc_history.append(epoch_acc) # 记录验证准确率
                if epoch_acc > best_acc:
                    best_acc = epoch_acc
                    best_model_wts = copy.deepcopy(model.state_dict())
                    print(f'  >> 新的最佳验证准确率: {best_acc:.4f}')

        print()

    time_elapsed = time.time() - since
    print(f'训练完成,耗时 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'最佳验证集准确率 (Best Val Acc): {best_acc:4f}')

    model.load_state_dict(best_model_wts)
    return model, val_acc_history

# -----------------------------------
# --- 5. 开始训练 ---
# -----------------------------------
# 设置训练轮数
num_training_epochs = 15 # 可根据需要调整

# 训练!
model_ft_trained, val_acc_hist = train_model(model_ft, criterion, optimizer_ft,
                                            # exp_lr_scheduler, # 如果使用调度器
                                             num_epochs=num_training_epochs)

# -----------------------------------
# --- 6. 保存模型 ---
# -----------------------------------
model_save_path = 'resnet18_finetuned_catdog_final.pth'
torch.save(model_ft_trained.state_dict(), model_save_path)
print(f"\n训练好的模型已保存至: {model_save_path}")

# (可选) 可视化训练过程中的验证准确率
# import matplotlib.pyplot as plt
# plt.figure()
# plt.title("Validation Accuracy vs. Number of Training Epochs")
# plt.xlabel("Training Epochs")
# plt.ylabel("Validation Accuracy")
# # 需要将 Tensor 类型的准确率转换为 Python float
# plt.plot(range(1, num_training_epochs+1), [acc.cpu().numpy() for acc in val_acc_hist])
# plt.ylim((0,1.))
# plt.xticks(range(1, num_training_epochs+1))
# plt.legend()
# plt.show()

请务必将代码中的 /path/to/your/cat_dog_dataset 替换为你自己存放猫狗数据集的实际路径!

四、 常见问题与进阶技巧

4.1 如何选择合适的预训练模型?

  • 任务相似度: 目标任务与预训练任务(通常是 ImageNet 分类)越相似,迁移效果通常越好。猫狗分类与 ImageNet 的物体识别任务相关度较高。
  • 数据集大小: 目标数据集非常小时,选择层数较少、参数量较小的模型(如 ResNet18/34)并冻结大部分层可能更稳妥,防止过拟合。数据集稍大时,可以尝试更深的模型(ResNet50+)或微调更多层。
  • 计算资源: 更深、更宽的模型(如 ResNet152, EfficientNet-B7)性能可能更好,但也需要更多的显存和计算时间。根据你的硬件条件选择。
  • 模型特性: 不同的模型架构各有侧重,例如 EfficientNet 在效率(性能与计算量平衡)上做了优化。可以查阅相关论文或基准测试结果进行选择。

4.2 如何设置学习率?

  • 基础学习率: 微调时通常使用比从零训练更小的学习率,例如 1e-31e-4
  • 差异化学习率 (Differential Learning Rates): 一个常见的技巧是为模型的不同部分设置不同的学习率。例如,对新添加的分类层使用较大的学习率(因为它需要从随机初始化开始学习),而对预训练的、被解冻的卷积层使用非常小的学习率(因为它们只需要微小的调整)。在 PyTorch 中,可以在定义优化器时传入多个参数组,每个组指定不同的学习率。
    # 示例:差异化学习率
    fc_params = list(map(id, model_ft.fc.parameters()))
    base_params = filter(lambda p: id(p) not in fc_params, model_ft.parameters())
    optimizer = optim.Adam([
        {'params': base_params, 'lr': 1e-4}, # 预训练层使用较小学习率
        {'params': model_ft.fc.parameters(), 'lr': 1e-3} # 新分类层使用较大学习率
    ], lr=1e-3) # 默认学习率(虽然这里被覆盖了)
    
  • 学习率调度器 (Learning Rate Scheduler): 在训练过程中动态调整学习率,如 StepLR(按步长衰减)、ReduceLROnPlateau(当指标不再提升时衰减)等,有助于模型更好地收敛和跳出局部最优。

4.3 过拟合与欠拟合怎么办?

  • 过拟合 (Overfitting): 模型在训练集上表现很好,但在验证集上表现差。
    • 解决方法:
      • 增加数据: 获取更多标注数据或使用更强的数据增强技术 (e.g., transforms.ColorJitter, transforms.RandomRotation, Mixup, CutMix)。
      • 减少模型复杂度: 使用更浅的网络(如 ResNet18 而非 ResNet50),或者冻结更多的层。
      • 正则化: 在优化器中增加 weight_decay(L2 正则化),或者使用 Dropout (虽然在预训练模型微调时,有时效果不明显或需要谨慎使用)。
      • 早停 (Early Stopping): 监控验证集性能,在性能不再提升时停止训练。我们的 train_model 函数实际上已经实现了保存最佳模型权重的逻辑,这也是一种形式的早停。
  • 欠拟合 (Underfitting): 模型在训练集和验证集上表现都不好。
    • 解决方法:
      • 增加模型复杂度: 使用更深或更宽的网络,或者解冻/微调更多的预训练层。
      • 训练更长时间: 确保模型有足够的时间学习。
      • 调整学习率: 学习率可能过小,尝试增大一点或使用更有效的学习率策略。
      • 检查数据预处理: 确保数据加载和预处理步骤正确无误。
      • 换用更好的预训练模型: 当前模型可能不足以捕捉数据的复杂性。

4.4 进一步提升性能?

  • 更强的预训练模型: 尝试 VGG, Inception, DenseNet, EfficientNet, ViT 等其他预训练模型。
  • 更丰富的数据增强: 探索 AutoAugment, RandAugment, Mixup, CutMix 等高级数据增强策略。
  • 集成学习 (Ensemble Methods): 训练多个不同的模型(或同一模型使用不同初始化/数据子集训练),并将它们的预测结果进行融合(如投票或平均)。
  • 优化器选择: 尝试不同的优化器,如 SGD with momentum, AdamW 等,并仔细调整其超参数。
  • 超参数调优: 使用网格搜索、随机搜索或贝叶斯优化等方法系统地寻找最佳的超参数组合(学习率、batch size、冻结策略、优化器参数等)。
  • 测试时增强 (Test Time Augmentation, TTA): 在预测阶段,对输入图像进行多种增强(如翻转、不同裁剪),分别预测后综合结果,可以提高预测的鲁棒性。

五、 总结

本文详细介绍了迁移学习的核心概念、优势及其在 PyTorch 中的实战应用。我们通过一个经典的猫狗分类任务,演示了如何利用强大的预训练模型 ResNet 进行微调,从而高效地构建高性能图像分类器。核心步骤总结如下:

  1. 理解迁移学习: 认识到其价值在于利用已有知识加速学习、减少数据依赖并提升模型性能。
  2. 选择预训练模型: 了解 ResNet 等常用模型的特点,并根据任务需求选择合适的模型(如 ResNet18)。
  3. 数据准备: 使用 torchvision.transforms 进行数据预处理和增强,关键是使用与预训练模型一致的归一化参数。利用 ImageFolderDataLoader 高效加载数据。
  4. 模型加载与修改: 使用 torchvision.models 加载预训练权重 (pretrained=True),并**替换掉模型的最后一层(分类头)**以适应目标任务的类别数量。
  5. 微调策略: (可选但常用) 冻结部分或全部预训练层的参数 (requires_grad=False),只训练新添加的层或少量解冻的层,以保留通用特征并防止过拟合。
  6. 定义训练组件: 选择合适的损失函数(如 nn.CrossEntropyLoss)和优化器(如 Adam, SGD),注意优化器只应更新 requires_grad=True 的参数,并通常设置较小的学习率
  7. 训练与验证: 编写标准的训练循环,包含前向传播、损失计算、反向传播和参数更新。在验证集上评估模型性能,并保存表现最佳的模型权重
  8. 问题排查与进阶: 了解如何根据训练/验证曲线判断过拟合或欠拟合,并掌握相应的解决策略,以及进一步提升性能的技巧(如差异化学习率、更强模型、高级增强等)。

通过掌握迁移学习,你可以显著提升在各种计算机视觉任务(甚至扩展到 NLP 等领域)中的开发效率和模型效果,尤其是在数据有限的情况下。希望本文能为你打开迁移学习的大门,并在你的 PyTorch 实战之路上提供有力的支持!


你可能感兴趣的:(PyTorch,pytorch,迁移学习,分类,ResNet,猫狗分类,人工智能,深度学习)