PyTorch实战：5分钟搞定torchvision.models预训练模型加载与微调（附代码）

PyTorch实战：5分钟搞定torchvision.models预训练模型加载与微调（附代码）

如果你刚开始接触PyTorch，面对一个全新的图像分类任务，是不是常常感到无从下手？自己从头搭建一个卷积神经网络，不仅需要深厚的理论基础，还得花费大量时间调试参数，结果可能还不尽如人意。别担心，PyTorch社区早就为你准备好了“开箱即用”的解决方案——torchvision.models。这个模块就像是一个预训练好的模型库，里面存放着经过ImageNet大赛千锤百炼的经典网络结构，比如ResNet、VGG、DenseNet等。你不需要理解它们内部每一层卷积的细节，就能直接调用，并快速应用到自己的数据集上。这篇文章，我就带你绕过那些繁琐的理论，直接上手实战，用最短的时间，让你手中的模型“跑”起来，并学会如何“调教”它，让它为你自己的任务服务。我们会从最基础的模型加载开始，一步步深入到模型结构的修改、参数的微调，最后还会分享几个我实际项目中踩过的坑和解决技巧。无论你是想快速验证一个想法，还是为产品开发原型，这套流程都能帮你节省大量时间。

1. 初识torchvision.models：你的深度学习工具箱

torchvision是PyTorch生态中专门为计算机视觉任务设计的工具包，它和torch的关系，就像是Photoshop的滤镜插件和Photoshop软件本身。torchvision.models则是这个工具包里的“模型博物馆”，收藏了计算机视觉领域过去十年里最具影响力的架构。这些模型并非空壳，它们都附带了在ImageNet数据集上训练好的权重。ImageNet包含上千万张图片、上千个类别，能在这个数据集上表现出色的模型，其提取图像特征的能力已经非常强大。

这意味着什么？意味着你可以直接“借用”这些模型已经学到的、对于通用图像（如边缘、纹理、形状）的识别能力，而不必从零开始。在学术上，这被称为迁移学习。想象一下，一个已经学会识别猫、狗、汽车的人，再去学习识别特定的猫品种，肯定比一个对图像一无所知的人要快得多。我们的模型也是如此。

使用torchvision.models，你通常只需要一行代码就能请来一位“世界级专家”。例如，想要使用在ImageNet上准确率很高的ResNet-50模型，你只需要：

import torchvision.models as models

resnet50 = models.resnet50(pretrained=True)

执行这行代码时，PyTorch会自动从云端下载预训练好的权重文件。pretrained=True这个参数就是关键，它告诉程序：“我要那个已经训练好的模型，别给我一个随机初始化的空壳。”

注意：首次运行下载可能会花费一些时间，取决于你的网络速度和模型大小。模型文件通常会缓存在本地，下次加载就快了。

这些预训练模型对输入数据有特定的要求，这并非刁难，而是为了保证模型能发挥出最佳性能。主要要求如下：

• 输入格式：一个RGB图像的小批量（mini-batch），形状为 (batch_size, 3, H, W)。其中H和W代表图像的高和宽，通常不能小于224像素。
• 像素值范围：图像张量的值应该在 [0, 1] 之间（即经过除以255的归一化）。
• 标准化：需要使用特定的均值和标准差进行标准化处理。这个均值和标准差是模型在ImageNet上训练时使用的。

为了方便，torchvision.transforms提供了现成的组合来处理这些要求：

from torchvision import transforms

# 定义标准的预处理流程
preprocess = transforms.Compose([
    transforms.Resize(256), # 将短边缩放到256
    transforms.CenterCrop(224), # 中心裁剪出224x224的区域
    transforms.ToTensor(), # 转换为张量，并归一化到[0,1]
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]) # 标准化
])

理解并正确使用这个预处理流程，是成功应用预训练模型的第一步。

2. 加载与窥探：让模型为你所用

加载模型只是第一步，我们还需要了解它的结构，知道从哪里入手进行修改。PyTorch的模型都是nn.Module的子类，结构非常清晰。让我们以VGG16为例，看看如何加载并观察它。

import torch
import torchvision.models as models

# 加载预训练的VGG16模型
vgg16 = models.vgg16(pretrained=True)

# 将模型设置为评估模式（重要！）
vgg16.eval()

# 打印模型结构
print(vgg16)

运行print(vgg16)，你会看到一长串的结构输出。它主要分为两大部分：features和classifier。

• features（特征提取器）：这部分由一系列的卷积层、激活函数和池化层堆叠而成，负责从原始图像中提取出越来越抽象的特征。对于VGG16，features部分会输出一个512x7x7的特征图。
• classifier（分类器）：这部分通常是几个全连接层，负责将features提取出的高级特征映射到最终的类别概率上。在预训练模型中，classifier的最后一层输出维度是1000，对应ImageNet的1000个类别。

关键操作：查看与修改分类器

我们的任务往往不是1000类的ImageNet分类，比如可能是10类的CIFAR-10，或是2类的猫狗分类。因此，修改classifier的最后一层是微调中最常见的操作。

方法一：直接替换最后一层 这是最直接、最常用的方法。我们找到classifier的最后一层（通常是最后一个Linear层），将其替换为一个输出维度为我们任务类别数的新层。

import torch.nn as nn

# 假设我们的新任务是10分类（如CIFAR-10）
num_ftrs = vgg16.classifier[6].in_features # 获取原最后一层的输入特征数
vgg16.classifier[6] = nn.Linear(num_ftrs, 10) # 替换为新的全连接层

print(vgg16.classifier) # 查看修改后的分类器

方法二：添加新的层 有时，我们可能不想完全丢弃预训练的分类器，而是想在它后面追加新的层。这可以通过add_module方法实现。

# 在原有分类器末尾添加一个新的全连接层
vgg16.classifier.add_module('7', nn.Linear(1000, 10)) # 添加一个名为‘7’的层

提示：方法一更常见，因为它直接让模型适应新任务的输出空间，且需要训练的新参数更少。方法二可能会引入更多的可训练参数，需要更谨慎地调整学习率。

修改完成后，模型的结构就适配了我们自己的任务。但此时，新添加或替换的层的参数是随机初始化的，而前面features和classifier中未修改层的参数仍然是预训练好的。接下来，我们需要决定如何训练这些参数。

3. 微调策略：如何高效训练你的模型

模型结构改好了，直接全部扔进优化器里训练吗？这通常不是最佳策略。预训练模型的前面部分已经学到了非常好的通用特征，我们应该以不同的“态度”对待它们。这里介绍两种核心的微调策略。

策略一：特征提取器冻结训练 这种策略将预训练好的features部分（特征提取器）完全冻结，不更新其参数。只训练我们新修改的classifier部分。这相当于把预训练模型当作一个固定的特征提取器，然后在其上训练一个简单的分类器（如SVM或一个小型神经网络）。这种方法训练速度快，计算资源消耗少，非常适合以下场景：

• 你的数据集与ImageNet差异较小。
• 你的数据集量很小，防止过拟合。
• 你只是想快速验证一个想法。

# 冻结 features 部分的所有参数
for param in vgg16.features.parameters():
    param.requires_grad = False

# 只优化 classifier 部分的参数
optimizer = torch.optim.SGD(vgg16.classifier.parameters(), lr=0.001, momentum=0.9)

策略二：整体微调 这种策略允许模型的所有参数都参与训练，但通常会对不同部分使用不同的学习率。我们给新添加的、随机初始化的层设置一个较大的学习率，让它快速学习；给预训练好的层设置一个较小的学习率，让它们进行精细的调整。这种方法潜力更大，通常能获得更高的精度，适用于：

• 你的数据集与ImageNet有一定差异，但仍有相似之处（如医学影像、卫星图片）。
• 你的数据集量足够大。
• 你对模型性能有较高要求。

实现不同层不同学习率，可以借助优化器的参数分组功能：

# 将模型参数分为两组
ignored_params = list(map(id, vgg16.classifier[6].parameters())) # 新层参数id
base_params = filter(lambda p: id(p) not in ignored_params, vgg16.parameters())

optimizer = torch.optim.SGD([
    {'params': base_params, 'lr': 0.001 * 0.1}, # 预训练层，小学习率
    {'params': vgg16.classifier[6].parameters(), 'lr': 0.001} # 新层，大学习率
], momentum=0.9)

为了更清晰地对比这两种策略，我们可以看看下面的表格：

在实际项目中，我通常会先尝试“特征提取”策略，快速得到一个基线模型。如果效果不理想，再切换到“整体微调”，并耐心调整学习率等超参数。记住，没有放之四海而皆准的策略，最佳选择取决于你的具体数据和任务目标。

4. 完整实战：从数据到训练的全流程

理论说再多，不如亲手跑一遍代码。让我们用一个完整的、可运行的例子，将CIFAR-10数据集在ResNet-18模型上进行微调。CIFAR-10是32x32的小图片，有10个类别，与ImageNet的224x224大图有所不同，正好可以展示预处理和微调的技巧。

第一步：准备数据 我们将使用torchvision.datasets和DataLoader来加载和批处理数据。

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 数据增强和预处理
# 注意：CIFAR-10图片尺寸为32x32，我们将其放大并裁剪，以适配ResNet输入
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.Resize(224), # 关键：上采样到224x224
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), # CIFAR-10的统计值
])

transform_test = transforms.Compose([
    transforms.Resize(224), # 关键：上采样到224x224
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

# 加载数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform_train)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform_test)
testloader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

第二步：准备模型 加载ResNet-18，并修改其最后的全连接层，以适应10分类任务。这里我们采用整体微调策略。

import torch.nn as nn
import torch.optim as optim

# 加载预训练模型
model = models.resnet18(pretrained=True)

# 修改最后一层全连接层
num_ftrs = model.fc.in_features # ResNet的最后一层是 `fc`
model.fc = nn.Linear(num_ftrs, 10) # CIFAR-10有10类

# 如果有GPU，将模型移至GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
# 对新层使用较高学习率，对其他层使用较低学习率
optimizer = optim.SGD([
    {'params': model.layer4.parameters(), 'lr': 0.001}, # 较深层，中等学习率
    {'params': model.fc.parameters(), 'lr': 0.01} # 新层，较高学习率
], momentum=0.9, weight_decay=5e-4)

第三步：训练与验证循环 这里展示一个简单的训练周期循环。

def train_one_epoch(model, trainloader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad() # 梯度清零
        outputs = model(inputs) # 前向传播
        loss = criterion(outputs, labels) # 计算损失
        loss.backward() # 反向传播
        optimizer.step() # 更新参数

        running_loss += loss.item()
    return running_loss / len(trainloader)

# 训练几个周期
for epoch in range(5):
    train_loss = train_one_epoch(model, trainloader, criterion, optimizer, device)
    print(f'Epoch {epoch+1}, Loss: {train_loss:.3f}')

这个流程为你提供了一个坚实的起点。你可以通过增加训练周期、添加学习率调度器（如StepLR）、在验证集上评估精度等方式，不断完善你的训练 pipeline。

5. 避坑指南与高阶技巧

在实际操作中，你肯定会遇到各种各样的问题。我把自己和同事们踩过的几个典型“坑”总结在这里，希望能帮你节省时间。

坑一：忘记设置 model.eval() 和 model.train() 这是一个非常常见但后果严重的错误。

• model.train()：在训练时调用，会启用Dropout和BatchNorm层的训练模式。
• model.eval()：在验证或测试时调用，会固定Dropout和BatchNorm层的运行统计量。 如果搞混，在验证时Dropout依然起作用，或者BatchNorm使用批统计量，会导致结果极不稳定且性能下降。

坑二：输入尺寸不匹配 预训练模型通常要求输入尺寸至少为224x224。如果你的图片是32x32（如CIFAR-10）或更大尺寸，直接输入会报错。务必在数据预处理中使用Resize或CenterCrop 来调整尺寸。对于小图片上采样，可能会损失一些信息，但这是使用标准预训练模型必须做的妥协。

坑三：预处理标准化参数用错 前面提到，ImageNet预训练模型要求使用mean=[0.485, 0.456, 0.406]和std=[0.229, 0.224, 0.225]进行标准化。如果你用自己的数据集从头训练，或者使用其他预训练模型（如在CIFAR上训练的模型），就需要使用对应的统计量。用错标准化参数相当于给模型提供了“失真”的输入，会严重影响性能。

高阶技巧：学习率热身与余弦退火 对于微调，特别是整体微调，学习率的设置非常关键。一个不错的实践是使用学习率热身：在训练的最开始几个step或epoch，让学习率从一个很小的值线性增长到预设的初始值，这有助于稳定训练初期。 之后，可以采用余弦退火策略来衰减学习率，让学习率随着训练过程平滑地降低，有助于模型收敛到更好的局部最优点。

from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR

# 假设optimizer已经定义
# 先热身（例如，用1个epoch热身）
warmup_epochs = 1
total_epochs = 50
warmup_scheduler = LinearLR(optimizer, start_factor=0.01, total_iters=len(trainloader)*warmup_epochs)
# 再余弦退火
cosine_scheduler = CosineAnnealingLR(optimizer, T_max=(total_epochs - warmup_epochs) * len(trainloader))

# 在每个batch的优化器step之后调用
for epoch in range(total_epochs):
    for batch in trainloader:
        # ... 训练步骤 ...
        optimizer.step()
        if epoch < warmup_epochs:
            warmup_scheduler.step() # 热身阶段
        else:
            cosine_scheduler.step() # 余弦退火阶段

模型保存与加载的注意事项 训练好的模型需要保存下来。PyTorch推荐只保存模型的state_dict()（状态字典），而不是整个模型对象。这样更灵活，且与文件大小无关。

# 保存
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, 'checkpoint.pth')

# 加载
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']

这种保存方式包含了训练状态，方便你从中断的地方继续训练。记住，加载时，你需要先实例化一个结构相同的模型和优化器，然后再加载状态字典。