【深度学习项目实战·15】小样本图像分类：迁移学习识别蚂蚁与蜜蜂

文章目录

- [01 本篇要解决什么问题](#01 本篇要解决什么问题)
- [02 背景：预训练参数提供一个更有信息的起点](#02 背景：预训练参数提供一个更有信息的起点)
- [03 实操：把预训练模型接到自己的小数据集](#03 实操：把预训练模型接到自己的小数据集)
- [04 一个具体例子：从目录数据训练到单图预测](#04 一个具体例子：从目录数据训练到单图预测)
- [05 常见问题与自测](#05 常见问题与自测)
- [06 本篇回顾与下一步](#06 本篇回顾与下一步)

01 本篇要解决什么问题

前几个项目都有数万张训练图。如果手里只有几百张彩色照片，从随机参数开始训练一个大网络，数据和训练预算可能都不够。

本篇使用预训练ResNet18，把已有图像特征迁移到蚂蚁与蜜蜂分类。重点区分只训练分类头与完整微调，并让模型的冻结状态和预处理真正对应起来。

本篇关键词：迁移学习、预训练、冻结参数、微调、BatchNorm。
完成标志：在小样本目录数据上训练分类器，保存类别映射，并能重新加载预测真实彩色照片。
运行约定 ：使用完整配套examples目录。首次运行会分别下载数据和预训练权重，两者不是同一个文件。

02 背景：预训练参数提供一个更有信息的起点

迁移学习利用源任务学到的表示，为新任务提供初始化或固定特征。它能减少新任务从头学习的负担，但是否有效仍取决于图像域、任务、预处理和训练方式。

本篇使用PyTorch官方迁移学习教程中的Hymenoptera小数据包，识别蚂蚁与蜜蜂，下载入口见官方教程。本包通常包含244张原始训练图与153张原始验证图，共397张。

为了保持本专栏训练、验证、最终测试三部分的边界，我们从原始train中再按类别划出20%用于验证，把原始val保留为最终留出集。按本版规则得到195张训练、49张验证、153张最终测试。

这里的"测试"是本篇对原有目录用途的重新约定，并不是数据发布方额外提供了一个官方测试目录。

03 实操：把预训练模型接到自己的小数据集

（1）下载后先检查目录和类别

powershell 复制代码

python examples/download_data.py --dataset transfer --data data

解压后目录大致为：

text 复制代码

data/hymenoptera_data/
├─ train/
│  ├─ ants/
│  └─ bees/
└─ val/
   ├─ ants/
   └─ bees/

ImageFolder按子目录组织类别，类别到数字的映射由数据集对象提供。代码会检查两个来源目录映射一致，并随权重保存，避免推理时把类别顺序写反。

（2）输入处理要与预训练权重匹配

使用ResNet18_Weights.DEFAULT加载ImageNet预训练参数。验证与推理预处理通过该权重的transforms()取得，包含匹配的缩放、中心裁剪和标准化配置，见ResNet18权重说明。

训练使用224×224随机裁剪与水平翻转，标准化沿用权重约定。不要直接把第14篇CIFAR的32×32与0.5均值标准差原封不动搬过来。

与从头训练小图网络不同，本篇保留标准ResNet18入口，替换最后的分类层为两类输出。

（3）先只训练分类头

冻结主干参数，只让新建的全连接层学习，是一种清楚的起点。它降低了需要优化的参数数量，也便于观察预训练特征对这个小任务是否有用。

下面是本篇使用的模块：

python 复制代码

from torch import nn
from torchvision.models import resnet18, ResNet18_Weights

class TransferClassifier(nn.Module):
    def __init__(self, mode="head", pretrained=True):
        super().__init__()
        self.mode = mode
        weights = ResNet18_Weights.DEFAULT if pretrained else None
        self.backbone = resnet18(weights=weights)
        self.backbone.fc = nn.Linear(self.backbone.fc.in_features, 2)
        if mode == "head":
            self.backbone.requires_grad_(False)
            self.backbone.fc.requires_grad_(True)

    def train(self, mode=True):
        super().train(mode)
        if self.mode == "head":
            # Frozen BatchNorm running statistics must remain frozen too.
            self.backbone.eval()
            self.backbone.fc.train(mode)
        return self

    def forward(self, x):
        return self.backbone(x)

requires_grad=False控制是否计算参数梯度；优化器还需要只接收打算更新的参数。二者都要检查，避免配置与实际更新范围不一致。

（4）冻结参数后，别忘了BatchNorm状态

BatchNorm的运行均值和方差是缓冲区，不是普通可训练权重。即使参数已经冻结，在训练状态下运行，统计量仍可能更新，相关行为见BatchNorm2d官方说明。

本篇覆盖模型的train()方法：只训练分类头时，主干始终保持评估状态，分类头按需要切换。这样不会被公共训练循环每轮调用model.train()意外解除。

验证与推理仍同时使用eval()和no_grad()。冻结参数、模块状态和梯度记录是不同控制维度，需要分别明确。

（5）再用较小学习率微调

--mode finetune解冻整个网络，主干学习率设为分类头的十分之一。默认分类头为0.001，主干为0.0001，这是一组教学起点，而非最优保证。

完整微调时BatchNorm随模型进入训练状态。小样本、较小batch可能使统计量不稳定，是否继续固定BN可以成为后续实验变量，但应明确记录，不能把两种配置混着比较。

本版head和finetune各自从同一预训练来源开始，不会自动把前一次head运行的权重作为finetune初始化。如果研究两阶段连续训练，应另外显式加载第一阶段权重并记录实验流程。

（6）用同一划分比较两种策略

powershell 复制代码

python examples/train_transfer.py --mode head --epochs 5 --output artifacts/15_head
python examples/train_transfer.py --mode finetune --epochs 5 --output artifacts/15_finetune

两次默认seed相同，因此数据划分一致。按验证结果分析策略，不要看一次测试表现就反复改变训练设置。最终留出集很小，少量错分就会明显改变准确率，解释结果时应同时给出样本数和错分图。

--no-pretrained只用于断网时的接口检查，它会从随机参数初始化，不能把这种运行结果称为迁移学习效果。

图：本篇原创示意。只训练分类头时，主干参数与BatchNorm运行统计均保持固定。

04 一个具体例子：从目录数据训练到单图预测

完整主入口为examples/train_transfer.py，所需数据和模型公共文件随工程提供，并收录于公共代码附录。

python 复制代码

"""Chapter 15: feature extraction or full fine-tuning from ImageNet weights."""
import torch
from project_common import parser_for, validate_args, seed_everything, run_experiment
from project_data import load_data, preprocessing
from project_models import build_model


def main():
    parser = parser_for(__doc__,"artifacts/15_transfer",epochs=5,batch_size=16)
    parser.add_argument("--mode",choices=["head","finetune"],default="head")
    parser.add_argument("--no-pretrained",action="store_true",help="For offline interface checks, not transfer learning")
    args = parser.parse_args()
    validate_args(parser,args)
    seed_everything(args.seed,args.threads)
    data,classes,manifest = load_data("transfer",args)
    spec = {"kind":"transfer_resnet18","mode":args.mode}
    model = build_model(spec,pretrained=not args.no_pretrained).to(args.device)
    if args.mode == "head":
        optimizer = torch.optim.Adam(model.backbone.fc.parameters(),lr=args.lr)
    else:
        backbone = [p for n,p in model.backbone.named_parameters() if not n.startswith("fc.")]
        optimizer = torch.optim.Adam([
            {"params":backbone,"lr":args.lr*0.1},
            {"params":model.backbone.fc.parameters(),"lr":args.lr}])
    run_experiment(model,spec,data,classes,manifest,preprocessing("transfer"),args,optimizer=optimizer)


if __name__ == "__main__":
    main()

小规模检查命令：

powershell 复制代码

python examples/train_transfer.py --mode head --epochs 1 --limit-train 32 --limit-val 16 --limit-test 16
python examples/visualize_run.py --run artifacts/15_transfer

把一张待预测图片保存为data/my_insect.jpg，执行：

powershell 复制代码

python examples/predict_image.py --checkpoint artifacts/15_transfer/best.pt --image data/my_insect.jpg

推理代码先用weights=None重建结构，再读取已经保存的完整权重，因此不需要重新下载ImageNet权重。预测分数仅对应蚂蚁和蜜蜂两个已知类别；输入其他昆虫时，模型仍会给出这两类中的结果，不能由此认定它具有未知类别识别能力。

05 常见问题与自测

问题1：数据很少时，迁移学习一定更好吗？

没有这种保证。先与合理起点比较，再看源域、新域和训练设置是否匹配。本篇提供的是可以验证的实现路线。

问题2：参数已经冻结，为什么模型状态仍会变？

检查BatchNorm运行统计等缓冲区。冻结梯度不会自动把这些模块切到评估状态。

问题3：两类任务为什么还能用原来1000类模型？

使用其预训练主干作为特征起点，最后1000类分类头替换为新建的两类头，再学习新类别。

问题4：head之后运行finetune，是接着训练吗？

本版命令不是。两次各自从预训练初始化开始，输出目录也应分开。连续两阶段训练需要明确加载上一阶段权重。

问题5：小留出集准确率变化很大，是不是代码错了？

先查流程，再看样本量。少量样本的错分就可能改变几个百分点，不能只凭一次数值判断策略稳定性。

06 本篇回顾与下一步

迁移学习需要同时处理权重来源、输入规范、可训练参数、模块状态与数据划分。把这些条件保存清楚，后面迁移到自己的图片目录会容易很多。

中文	常见英文或写法
特征提取	Feature Extraction
微调	Fine-tuning
冻结参数	Frozen Parameters
运行统计量	Running Statistics

下一篇进入医学图像二分类，重点处理类别不平衡、验证集阈值和多指标评价。

参考与延伸阅读