1. 为什么迁移学习是图像分类任务的第一选择
如果你刚开始接触深度学习图像分类,大概率会遇到两个绕不开的坎:一是自己从头搭一个卷积网络,训练半天准确率上不去;二是手头的数据集太小,几千张图片都算多的,模型还没怎么学就过拟合了。这两个问题其实指向同一个解决方案——迁移学习。
迁移学习的核心思路很朴素:别人已经在大规模数据集上训练好的网络,已经学会了如何提取图像的通用特征,比如边缘、纹理、形状、局部结构。这些特征对绝大多数图像任务都是通用的,你不需要让网络从零开始理解什么是“边缘”,只需要让它学会“这个任务的边缘组合长什么样”就够了。VGG16 就是在 ImageNet 上预训练过的经典网络,结构规整、特征提取能力强,非常适合作为迁移学习的起点。
这篇文章面向的是有一定 PyTorch 基础、想快速跑通一个图像分类项目的读者。我会从数据准备、模型改造、训练策略、踩坑排查几个维度,把 VGG16 迁移学习的完整流程拆开讲清楚。不是那种“跑通就行”的教程,而是把每一步为什么这么做、不这么做会出什么问题都说明白。读完你拿到一个新数据集,能自己判断该冻结哪些层、该用什么学习率、该不该做数据增强。
2. 迁移学习到底迁移了什么:VGG16 的结构拆解
2.1 VGG16 的卷积基与分类头分工
VGG16 的结构非常规整,13 个卷积层加 3 个全连接层,卷积部分全部使用 3×3 小卷积核,池化统一用 2×2 最大池化。这种设计的精妙之处在于:堆叠小卷积核可以在保持感受野的同时减少参数量,而且每一层后面都跟 ReLU,非线性表达能力足够强。
从迁移学习的角度看,VGG16 可以拆成两部分:
- 卷积基(features):从输入图像到最后一个池化层,输出 512×7×7 的特征图。这部分负责提取图像的视觉特征,是迁移学习的核心。
- 分类头(classifier):三个全连接层,最后输出 1000 维的 ImageNet 类别概率。这部分是针对 ImageNet 的 1000 类专门训练的,对你的任务没有直接用处,必须替换。
很多人第一次做迁移学习时会犯一个错误:直接把 VGG16 的 classifier 最后一层改成自己的类别数就开训。这样做不是不行,但效果往往不如预期,因为中间那两个 4096 维的全连接层参数量巨大,小数据集上极容易过拟合。
2.2 浅层特征与深层特征的差异
卷积基内部也不是铁板一块。浅层卷积(前几个 block)学到的是边缘、颜色、简单纹理这类低级特征,这些特征几乎对所有图像任务都通用。深层卷积(后面的 block)学到的是更抽象、更语义化的特征,比如“狗的眼睛”“车轮的形状”,这些特征和具体任务的相关性更强。
这就引出一个关键决策:到底冻结多少层?
我的经验是:
| 数据量 | 与 ImageNet 相似度 | 推荐策略 |
|---|---|---|
| 很小(<1000) | 高 | 只训练分类头,冻结全部卷积基 |
| 很小(<1000) | 低 | 冻结前几个 block,训练后几个 block + 分类头 |
| 中等(1000-10000) | 高 | 冻结前两个 block,训练其余部分 |
| 较大(>10000) | 任意 | 全部微调,但用较小学习率 |
这个表不是死规矩,但能帮你避开“一上来就全量微调”这个最常见的坑。小数据集全量微调,基本等于让一个博士生去做小学题,他会的那些高级知识反而会干扰判断。
2.3 为什么 VGG16 仍然值得学
现在新网络层出不穷,EfficientNet、ConvNeXt、ViT 在各项指标上都超过了 VGG16。但 VGG16 依然是学习迁移学习的最佳起点,原因有三:
第一,结构简单直观。没有残差连接、没有分组卷积、没有注意力机制,就是卷积-池化-全连接的线性堆叠,你很容易看清楚每一层在做什么。
第二,特征提取能力扎实。VGG16 在 ImageNet 上的 top-5 准确率超过 90%,提取的特征在迁移到其他任务时表现稳定,不会出现某些新网络“预训练权重迁移效果反而差”的情况。
第三,计算量可控。VGG16 参数量约 1.38 亿,其中绝大部分集中在全连接层。如果你只训练分类头,实际需要更新的参数只有几百万,普通显卡甚至 CPU 都能跑。
3. 数据准备:从文件夹到 DataLoader 的完整链路
3.1 目录结构决定加载方式
PyTorch 的ImageFolder要求数据按类别分文件夹存放,这是最省事的组织方式:
dataset/ ├── train/ │ ├── cat/ │ │ ├── 001.jpg │ │ └── ... │ └── dog/ │ ├── 001.jpg │ └── ... └── val/ ├── cat/ └── dog/这种结构下,ImageFolder会自动读取文件夹名作为类别标签,你不需要额外维护标签映射表。如果数据不是这个结构,要么手动整理,要么自己写 Dataset 子类。
注意:训练集和验证集的类别文件夹必须完全一致,否则验证时会出现类别索引错位,准确率直接崩掉。
3.2 数据增强的尺度怎么把握
数据增强是防止过拟合的第一道防线,但不是越多越好。VGG16 输入尺寸是 224×224,常用的增强手段包括:
- RandomResizedCrop(224):随机裁剪并缩放到 224,模拟不同拍摄距离和构图。
- RandomHorizontalFlip():水平翻转,适合大多数自然图像,但文字、数字识别任务要慎用。
- ColorJitter:调整亮度、对比度、饱和度,模拟不同光照条件。
- RandomRotation(15):小角度旋转,角度太大会引入不真实的样本。
验证集只做 Resize(256) + CenterCrop(224),不做任何随机增强,保证评估结果稳定可复现。
from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])这里的 mean 和 std 是 ImageNet 的统计值,必须和预训练权重保持一致。如果你自己算一套归一化参数,预训练权重的特征分布就对不上了,效果会明显下降。
3.3 DataLoader 的 batch_size 与 num_workers 取舍
batch_size 的选择和显存直接相关。VGG16 参数量大,即使只训练分类头,前向传播的中间激活也占不少显存。8GB 显存建议 batch_size 设为 16 或 32,先跑一个 batch 看看显存占用再调整。
num_workers 控制数据加载的并行进程数。设成 0 表示在主进程加载,调试时方便看报错;设成 4 或 8 能加速训练,但 Windows 上有时会遇到多进程问题,可以先从 0 开始,跑通后再调大。
from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset = ImageFolder('dataset/train', transform=train_transform) val_dataset = ImageFolder('dataset/val', transform=val_transform) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True)pin_memory=True在 GPU 训练时能加速 CPU 到 GPU 的数据传输,这个细节很多人会忽略,但确实能省一点时间。
4. 模型改造:冻结策略与分类头设计
4.1 加载预训练权重并替换分类头
import torch import torch.nn as nn from torchvision import models model = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1) # 冻结卷积基 for param in model.features.parameters(): param.requires_grad = False # 替换分类头 num_classes = len(train_dataset.classes) model.classifier = nn.Sequential( nn.Linear(512 * 7 * 7, 512), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(512, num_classes) )这里我把原来的三层全连接(4096-4096-1000)改成了两层(512-num_classes)。为什么这么改?因为原版的 4096 维全连接参数量太大,小数据集上训练不动还容易过拟合。降到 512 维,参数量减少了一个数量级,训练更快,泛化也更好。
Dropout 设 0.5 是经验值,如果你的数据集特别小(几百张),可以加到 0.6 甚至 0.7;如果数据量上万,可以降到 0.3。
4.2 冻结策略的渐进式解冻
如果你发现只训练分类头效果不够好,可以尝试渐进式解冻:先训练分类头几个 epoch,等 loss 稳定后,解冻最后一个卷积 block(features 的第 24-30 层),用更小的学习率继续训练。
# 先训练分类头 5 个 epoch # 然后解冻最后一个 block for param in model.features[24:].parameters(): param.requires_grad = True # 优化器要重新定义,因为参数的可训练状态变了 optimizer = torch.optim.Adam([ {'params': model.features[24:].parameters(), 'lr': 1e-5}, {'params': model.classifier.parameters(), 'lr': 1e-3} ])注意这里用了分组学习率:卷积层用 1e-5,分类头用 1e-3。因为卷积层加载的是预训练权重,已经比较好了,只需要微调;分类头是随机初始化的,需要更大的学习率快速收敛。
4.3 分类头设计的几种变体
除了上面那种两层结构,还有几种常见设计:
单层线性:nn.Linear(512*7*7, num_classes)。最简单,参数量最少,适合数据量极小的情况。
带 BatchNorm 的两层:在 Linear 后面加 BatchNorm1d,能加速收敛,但对 batch_size 敏感,batch_size 太小时效果不稳定。
全局平均池化 + 线性:把 512×7×7 的特征图做全局平均池化变成 512 维,再接 Linear。参数量大幅减少,但会丢失空间信息,适合分类任务不太依赖位置信息的场景。
我的建议是先用两层结构跑 baseline,如果过拟合明显再简化,如果欠拟合再加深。
5. 训练循环:损失函数、优化器与学习率调度
5.1 损失函数与优化器的选择逻辑
分类任务用 CrossEntropyLoss,这个没有太多争议。优化器方面,Adam 和 SGD 各有拥趸:
- Adam:自适应学习率,收敛快,对初始学习率不敏感,适合快速实验。
- SGD + Momentum:泛化性能通常更好,但需要调学习率和 momentum,收敛慢。
我的习惯是:第一阶段用 Adam 快速验证流程是否跑通,第二阶段换 SGD 精调。如果时间紧,全程 Adam 也能拿到不错的结果。
criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.classifier.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)StepLR 每 7 个 epoch 把学习率乘以 0.1,这是一种简单有效的衰减策略。如果你用 CosineAnnealingLR,学习率会平滑下降,通常效果更好,但需要设对 T_max。
5.2 训练与验证的完整代码框架
def train_one_epoch(model, loader, criterion, optimizer, device): model.train() running_loss = 0.0 correct = 0 total = 0 for images, labels in loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() * images.size(0) _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() return running_loss / total, correct / total def validate(model, loader, criterion, device): model.eval() running_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for images, labels in loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) running_loss += loss.item() * images.size(0) _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() return running_loss / total, correct / total验证阶段一定要加torch.no_grad(),否则会构建计算图,显存占用飙升,而且完全没必要。
5.3 学习率调度的时机与幅度
学习率调度不是越复杂越好。我见过有人在 10 个 epoch 的训练里用了三种调度策略,结果学习率忽高忽低,loss 曲线像心电图。简单任务用 StepLR 就够了,复杂任务可以试 CosineAnnealing 或 ReduceLROnPlateau。
ReduceLROnPlateau 的逻辑是:验证 loss 连续几个 epoch 不下降就降低学习率。这个策略比较“智能”,但 patience 参数要设对,设太小会频繁降学习率导致训练停滞。
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=3 ) # 每个 epoch 后调用 scheduler.step(val_loss)6. 踩坑实录:迁移学习中最容易翻车的六个地方
6.1 归一化参数不匹配导致特征分布偏移
这是最隐蔽的坑。你用了预训练权重,但归一化用的是自己数据集的 mean 和 std,结果预训练学到的特征分布完全对不上,准确率可能比随机猜还低。记住:预训练权重用什么归一化参数,你就用什么。ImageNet 预训练就用 ImageNet 的 mean 和 std,不要自作主张。
6.2 冻结层没生效:requires_grad 的陷阱
# 错误做法 for param in model.features.parameters(): param.requires_grad = False # 但后面又写了这个 optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)model.parameters()会返回所有参数,包括那些requires_grad=False的。虽然 PyTorch 会自动忽略不需要梯度的参数,但如果你不小心把冻结层的参数也传进去了,优化器状态会多占显存。正确做法是只传需要训练的参数:
optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3 )6.3 验证集准确率震荡大:batch_size 与 BatchNorm 的冲突
如果你在分类头里加了 BatchNorm,验证时 batch_size 又设得和训练时不一样,准确率会明显波动。因为 BatchNorm 在训练时用当前 batch 的统计量,验证时用 running_mean 和 running_var,batch_size 变化会影响统计量的稳定性。解决办法要么去掉 BatchNorm,要么保证训练和验证的 batch_size 一致。
6.4 显存溢出:VGG16 的中间激活占用
VGG16 的第一个全连接层输入是 512×7×7=25088 维,即使你改成了 512 维输出,中间激活也不小。如果显存不够,可以:
- 减小 batch_size
- 用
torch.cuda.amp混合精度训练 - 把模型放到 CPU 上跑(慢但能跑)
混合精度训练能省大约 40% 显存,而且速度更快,值得一试:
from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() with autocast(): outputs = model(images) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()6.5 过拟合的早期信号与应对
训练 loss 持续下降但验证 loss 开始上升,这是过拟合的典型信号。应对手段按优先级排序:
- 增加数据增强强度
- 增大 Dropout
- 减小分类头参数量
- 提前停止(Early Stopping)
- 加 L2 正则化(weight_decay)
weight_decay 在优化器里设置,Adam 默认是 0,可以设成 1e-4 或 1e-5。注意 Adam 的 weight_decay 实现和 SGD 不同,如果追求严格的正则化效果,可以用 AdamW。
6.6 类别不平衡:准确率虚高的假象
如果验证集里 90% 是猫,10% 是狗,模型全预测猫也能拿到 90% 准确率。这时候要看混淆矩阵和每类 F1 分数,不能只看总体准确率。处理方法包括:
- 重采样:对少数类过采样,对多数类欠采样
- 加权损失:给少数类更高的 loss 权重
- 换评估指标:用 macro-F1 代替 accuracy
# 加权 CrossEntropyLoss class_counts = [len(os.listdir(f'dataset/train/{c}')) for c in train_dataset.classes] weights = 1.0 / torch.tensor(class_counts, dtype=torch.float) weights = weights / weights.sum() criterion = nn.CrossEntropyLoss(weight=weights.to(device))7. 训练之后的模型评估与推理部署
7.1 混淆矩阵与分类报告
训练完成后,别只看一个准确率数字。用 sklearn 的 classification_report 和 confusion_matrix 看看每类的表现:
from sklearn.metrics import classification_report, confusion_matrix model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for images, labels in val_loader: images = images.to(device) outputs = model(images) _, predicted = outputs.max(1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_names=val_dataset.classes)) print(confusion_matrix(all_labels, all_preds))如果某一类 recall 特别低,说明模型对这一类不敏感,可能需要补充这类样本或调整损失权重。
7.2 单张图片推理的完整流程
from PIL import Image def predict(image_path, model, transform, classes, device): model.eval() image = Image.open(image_path).convert('RGB') image = transform(image).unsqueeze(0).to(device) with torch.no_grad(): outputs = model(image) probabilities = torch.softmax(outputs, dim=1) confidence, predicted = probabilities.max(1) return classes[predicted.item()], confidence.item()注意unsqueeze(0)是必须的,因为模型期望的输入是(batch_size, channels, height, width),单张图片只有三个维度,需要补一个 batch 维度。
7.3 模型保存与加载的两种方式
方式一:保存整个模型
torch.save(model, 'model.pth') model = torch.load('model.pth')简单粗暴,但文件大,而且依赖模型定义代码。
方式二:只保存参数(推荐)
torch.save(model.state_dict(), 'model_weights.pth') # 加载时需要先实例化模型结构 model = models.vgg16(weights=None) model.classifier = nn.Sequential(...) # 和训练时一致 model.load_state_dict(torch.load('model_weights.pth'))推荐方式二,文件小,灵活,而且不会因为代码重构导致加载失败。
8. 从 VGG16 出发的进阶方向
跑通 VGG16 迁移学习之后,你可以沿着几个方向继续深入。第一个方向是换 backbone,试试 ResNet50 或 EfficientNet-B0,感受一下残差连接和复合缩放带来的差异,迁移学习的流程完全一样,只是模型结构变了。第二个方向是尝试更精细的微调策略,比如逐层解冻、判别式学习率(不同层用不同学习率),这些技巧在 Kaggle 比赛里很常见。第三个方向是处理更复杂的场景,比如多标签分类、细粒度分类、小样本学习,这些任务对特征提取和分类头设计的要求更高。
我个人在实际操作中的体会是:迁移学习的效果上限很大程度上取决于预训练数据集和你目标领域的相似度。如果相似度高,简单冻结就能拿到很好的结果;如果相似度低,就需要更激进地解冻和微调。判断相似度的一个实用方法是:先用冻结方案跑几个 epoch,看验证准确率能不能超过 70%,能超过说明特征迁移有效,不能超过就要考虑解冻更多层或者换预训练权重。这个判断方法比盲目调参高效得多。