Swin-Transformer水果图像分类迁移学习实践指南

简介：面向图像分类与迁移学习实践者，这是一套基于Swin-Transformer的水果十二分类图像识别项目，可直接运行并支持替换为自己的数据集。数据集涵盖香蕉、苹果、西瓜等12类水果，包含2340张训练图片与581张预测图片；模型采用cos学习率自动衰减训练50个epoch，在测试集上最高精度达99.3%。包内共2000个文件，以jpeg图片数据为主，辅以png、webp等格式的扩充样本，另有4个Python脚本、json配置与readme说明，便于理解数据处理、训练及推理流程。压缩包大小约969.65MB。目前已有266人学习使用。该项目既适合深度学习入门者快速复现高精度分类流程，也为需要定制水果识别或迁移学习方案的研究者提供了完整可扩展的代码与数据基础。

1. 为什么用 Swin-Transformer 做水果十二分类迁移学习

水果数据集的十二分类任务，实际工程里往往卡在数据量上：光照不统一、腐烂样本干扰、单类图片可能只有几百张。从零训练 Swin-Transformer 不会比 ResNet-50 好；正确做法是加载 ImageNet-1K 预训练权重做微调，这也是深度学习图像识别项目里最稳妥的迁移学习路径。Swin 的分层多尺度特征和 CNN 的 stage 结构一致，适配成本低。

相对 ViT，Swin-Transformer 的移位窗口让注意力落在局部窗口内，收敛更稳，更适合小数据集微调。十二分类用 Swin-Tiny 就够，权重在 timm 里直接可用。

接下来按选型、代码、调参、评估的顺序展开，覆盖分类头替换、冻结策略、学习率和混合精度训练，以及验证模型学到的究竟是水果特征还是背景纹理。适合已经跑过 CNN 分类，准备换 Transformer 骨干的工程师。

2. Swin-Transformer 的结构拆解与迁移学习选型

2.1 Patch Embedding 与四阶段层级特征

Swin-Transformer 的输入处理方式和 ViT 类似，但 patch 更小：图像按 4×4 像素切块，每个 patch 展平成 48 维向量，经线性层投影为 C 维嵌入。随后网络依次经过 stage1 到 stage4，每个 stage 前由一个 Patch Merging 层负责将分辨率减半、通道数翻倍。Swin-Tiny 的初始通道 C=96，四个 stage 的通道序列是 96/192/384/768，对应输出特征图尺寸是 56×56/28×28/14×14/7×7。这和 ResNet-50 的 stage 输出尺寸完全对应，意味着分类头、FPN、注意力可视化等下游组件的接入方式可以直接沿用 CNN 时期的工程经验。

在图像识别模型的主流选择里，Swin-Transformer 是少有的兼顾精度和部署灵活性的分层 Transformer。对水果十二分类这个任务，分类头读的是 stage4 的 7×7×768 输出。做全局平均池化后接一个 768→12 的全连接层，模型参数量 2830 万个，绝大多数集中在 stage3 和 stage4，这也是后面微调时建议先解冻 stage4 的原因。

2.2 移位窗口注意力的计算量优势

标准 Transformer 的全局自注意力在 H=W=56 的特征图上的计算量很大，其中 HW×HW 项约 980 万对位置相互作用，而 Swin 的窗口注意力把特征图切成 8×8 个窗口、每个窗口 7×7，只需要约 15 万对，量级下降近 98.5%。这是 Swin 能在同样吞吐量下处理 224×224 输入的根基。

代价是窗口内看不到全图信息。Swin 的解法是交替使用 W-MSA 和 SW-MSA：规则窗口算完一次后，特征图向右下偏移 3 个 patch 重划窗口，两次注意力合起来覆盖相邻窗口的交互。这个设计也导致 Swin 和 ViT 在迁移行为上差异很大：Swin 对输入尺寸变化更敏感，因为相对位置编码是按窗口尺寸预训练的。

训练时显存主要由激活值而非参数决定。Swin-Tiny 在 batch 32、AMP 下显存约 12GB，显存吃紧时优先降 batch 而不是降输入分辨率，因为预训练权重和相对位置编码都绑定 224×224。窗口注意力把自注意力从 O(H²W²) 压到近似线性，这是实际部署里最直接的收益来源。

2.3 变体选择与预训练权重加载

用一张表决定选哪个变体：

对水果项目，我一般默认选 Swin-Tiny。Swin-Base 在 ImageNet-1K 上仅比 Tiny 高 2.2%，参数量却是 3 倍；迁移到 12 类任务后这个差距会更小，而显存压力是实打实的。如果确实遇到欠拟合，优先补数据而不是升模型。

timm 加载预训练权重和替换分类头的代码：

import timm model = timm.create_model( 'swin_tiny_patch4_window7_224', pretrained=True, num_classes=12 ) print(model.default_cfg)

代码说明：pretrained=True 时 timm 会下载 ImageNet-1K 权重，并把分类头替换为输出 12 类的线性层；default_cfg 里保存了该模型的预处理均值、标准差和输入尺寸，训练和推理时按这个配置做 Normalize，不要自行更换统计量。

直推式迁移学习是另一个分支，用于目标集几乎没有标签的场景，需要借助伪标签和领域自适应。本文的水果十二分类属于监督迁移，每个类别都有标注，不需要走直推式路线。

3. 数据准备与 Swin-Transformer 微调的代码实现

3.1 数据集目录组织与增强参数

我一般用四步准备水果数据集：收集图片、人工清洗、按类别放置目录、划分训练验证集。目录组织采用 ImageFolder 约定：

data/ ├── train/ │ ├── apple/ │ ├── banana/ │ ├── ... # 其余类别 └── val/ ├── apple/ ├── banana/ └── ...

训练增强和验证集预处理的代码：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(size=224, scale=(0.7, 1.0), ratio=(0.75, 1.333)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(degrees=15), 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]) ])

增强参数和迁移学习直接相关：RandomResizedCrop 的 scale 下限设 0.7 而不是 0.08，是因为水果通常在图中占比大，过度缩小会让模型只能靠背景判断类别。旋转限制在 15 度，避免出现大量无意义的倒置视角。验证集只做 Resize 和 CenterCrop，不做任何增强，否则验证结果不稳定。Normalize 用 ImageNet 统计量，因为预训练权重就是在该分布下收敛的。

3.2 用 timm 替换分类头的正确做法

替换分类头有两种写法，常见的是下面这组：

import timm import torch.nn as nn # 方式一：创建时直接指定类别数，内部完成 head 替换 model = timm.create_model( 'swin_tiny_patch4_window7_224', pretrained=True, num_classes=12 ) # 方式二：拿到模型后手动替换分类头 model.num_features # 768 model.head = nn.Linear(768, 12)

方式二容易踩坑：如果忘记确认 model.num_features，或把它和预训练权重的输出尺寸搞混，训练时会直接报形状不匹配。检查替换后的输出是否正确，可以跑一次随机输入：

import torch x = torch.randn(2, 3, 224, 224) y = model(x) assert y.shape == (2, 12)

这里的随机输入只是为了验证维度，不参与反向传播。实际训练时，backbone 参数保持预训练值，head 是随机初始化的，第一批数据反向传播会把随机初始化的 head 梯度传给 backbone，噪声很大——这就是为什么迁移学习要配 warmup，不能直接上大学习率。

3.3 DataLoader 配置与混合精度训练循环

DataLoader 的参数影响训练稳定性和显存效率，实际项目中的配置：

from torch.utils.data import DataLoader train_loader = DataLoader( train_dataset, batch_size=32, # 显存不够时降到 16，不要降到 8 以下 shuffle=True, num_workers=8, # Windows 建议 4，Linux 可 8~16 pin_memory=True, drop_last=True )

训练循环用 AMP 的写法：

import torch from torch.cuda.amp import autocast, GradScaler model = model.cuda() optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.05) criterion = torch.nn.CrossEntropyLoss(label_smoothing=0.1) scaler = torch.cuda.amp.GradScaler() def train_one_epoch(model, loader): model.train() total_loss = 0.0 for images, labels in loader: images = images.cuda() labels = labels.cuda() optimizer.zero_grad() with autocast(): logits = model(images) loss = criterion(logits, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) scaler.step(optimizer) scaler.update() total_loss += loss.item() * images.size(0) return total_loss / len(loader.dataset)

逻辑说明：autocast 让 forward 在 fp16 下计算，但 LayerNorm 和 softmax 会自动保持 fp32 精度。GradScaler 对 loss 做缩放，防止反向传播时梯度下溢为 0。clip_grad_norm 必须在 unscale_ 之后、step 之前执行，否则裁剪阈值作用在缩放后的梯度上会失效。label_smoothing 设为 0.1，对标注噪声较多的水果数据集能明显提高验证集稳定性，代价是训练 loss 不会再降到接近 0。

4. Swin-Transformer 迁移学习微调参数与收敛控制

4.1 分层学习率与优化器分组

微调 Swin 和微调 ResNet 的第一个区别就在学习率上。CNN 微调常用 5e-3 起步，Transformer 用这个值会直接让预训练特征崩溃。常见做法是 backbone 和 head 分开设置学习率，head 是随机初始化的，需要走更快的收敛路径。

backbone_params = [] head_params = [] for name, param in model.named_parameters(): if param.requires_grad: if name.startswith('head.'): head_params.append(param) else: backbone_params.append(param) optimizer = torch.optim.AdamW([ {'params': backbone_params, 'lr': 2e-5}, {'params': head_params, 'lr': 2e-4}, ], weight_decay=0.05)

head 学习率设为 backbone 的 10 倍是实践常用值。如果跑了 3 个 epoch 发现 loss 震荡，先除以 5，再观察；如果 head 收敛明显慢，提高倍数而不是提高 backbone 的学习率。weight_decay 用 0.05，这是 Swin 官方预训练使用的配置，和 CNN 常用的 1e-4 不是一个体系。

4.2 冻结与分阶段解冻的实操

当每类样本只有几百张时，冻结主干和分阶段解冻是非常有效的办法。冻结操作本身不复杂，关键是解冻哪个 stage：

# 阶段一：冻结全部 backbone，只训练 head for param in model.parameters(): param.requires_grad = False for param in model.head.parameters(): param.requires_grad = True # 阶段二：head 准确率稳定后，解冻最后一个 stage for name, param in model.named_parameters(): if name.startswith('layers.3.'): param.requires_grad = True

Swin-Tiny 的层名映射是 layers.0 到 layers.3，分别对应 stage1 到 stage4。解冻顺序从后往前，先解冻 layers.3，再根据验证集表现决定是否解冻 layers.2。冻结策略不是死规则，判断标准应该是任务规模和硬件资源：训练集总样本少于 6000 张时，先冻结主干只训 head，准确率稳定后再解冻最后一个 stage 做低学习率微调；训练集总量超过 2 万张时直接全量微调，并把 warmup 放在最前面。

4.3 Warmup、余弦退火与训练监控

Transformer 微调的调度器配置比 CNN 讲究，Warmup 长度直接决定前期稳定性。用 PyTorch 内置调度器组合实现：

from torch.optim.lr_scheduler import SequentialLR, LinearLR, CosineAnnealingLR warmup = LinearLR(optimizer, start_factor=0.001, end_factor=1.0, total_iters=3) cosine = CosineAnnealingLR(optimizer, T_max=27, eta_min=1e-6) scheduler = SequentialLR(optimizer, [warmup, cosine], milestones=[3])

每个 epoch 结束后调用 scheduler.step()。LinearLR 在 3 个 epoch 内从初始学习率的 0.1% 逐步升到 100%，SequentialLR 在第 3 个 epoch 结束后切换到余弦退火，剩下 27 个 epoch 平滑降到 1e-6。这里 T_max=27 是因为总 epoch 数减去 warmup 的 3 个 epoch。

迁移学习训练曲线有三种典型模式需要区分：train loss 持续下降但 val loss 上行，是过拟合，应加大增强或降低 backbone 学习率，而不是提前终止；train loss 一直降不下去，是 head 还没收敛就解冻了主干，回到阶段一再跑几个 epoch；loss 呈锯齿状剧烈波动，优先查 batch size 是否小于 16，以及数据加载是否混入坏样本。

5. 十二分类评估指标与模型导出验证

5.1 混淆矩阵与分类别指标

训练完成后，测试集上收集预测结果：

from sklearn.metrics import classification_report, confusion_matrix # preds: 所有测试样本的预测类别，labels: 真实类别 report = classification_report(labels, preds, target_names=class_names, digits=3) print(report) cm = confusion_matrix(labels, preds) cm_norm = cm.astype('float') / cm.sum(axis=1, keepdims=True) print(cm_norm.round(3))

重点看标准化混淆矩阵对角线以外的值是否大于 0.1。公开水果数据集中常有新鲜和腐烂样本并存的设置，腐烂果形态不规则，比新鲜果更容易误判；颜色相近的组合，比如苹果和梨或青番茄，也是高频混淆对。若某个类别召回率远低于总体准确率，问题不在模型容量，而在该类别样本数和采集条件多样性不足。

5.2 Grad-CAM 验证注意力是否落在水果上

迁移学习最隐蔽的失败是模型学到背景偏见，精度不低但根本没看水果。用 Grad-CAM 验证是成本最低的检查方式。Swin 的目标层选 model.layers.3[-1] 的 norm1 输出，把特征图梯度反传回输入，得到热力图后叠加到原图上。不同 timm 版本模块命名有差异，建议先打印 model 结构再选目标层。

如果热力图中聚焦区域集中在水果实体上，说明分类依据可靠；如果热力集中在叶片、篮筐或包装袋上，说明训练集里该类别总是伴随同一背景，模型在用上下文作弊。处理办法是收集该类在不同背景下的图片，或对该类做更强的随机裁剪增强，强迫模型关注外观。

5.3 ONNX 导出与推理一致性验证

模型验证通过后，导出到推理端常见做法是转 ONNX：

model.eval() dummy = torch.randn(1, 3, 224, 224, device='cuda') torch.onnx.export( model, dummy, "swin_fruit.onnx", input_names=["image"], output_names=["logits"], opset_version=16 )

用 onnxruntime 加载同一批测试图，比较 .pth 和 .onnx 的 argmax 结果，两者不一致通常来自浮点计算顺序差异，不是模型坏了。如果要求严格一致，导出时不声明 dynamic_axes，固定 224×224 输入。

onnxruntime 默认 CPU 执行，单张 224×224 的 Swin-Tiny 推理耗时大约 30 到 60ms；加上 CUDA Execution Provider 后可降到 5ms 以内。若延迟还不满足，最后再考虑转 TensorRT，但 Swin 的动态窗口在 TensorRT 部分版本下算子融合不完整，适配成本比 CNN 高，这是 Transformer 部署和卷积网络最大的区别之一。

本文还有配套的精品资源，点击获取