PyTorch训练CUB200全流程：细粒度图像分类数据准备与ResNet微调

简介面向CUB-200鸟类数据集的卷积神经网络分类训练代码包专为需要解决200类细粒度鸟类识别的学生与研究者设计涵盖了从数据读取到模型训练的关键环节。压缩包共含5个文件其中3个Python脚本分别承担数据预处理、模型结构定义与训练主流程另有2个pyc编译文件整体仅15KB代码精简且无需依赖大型工程即可运行。这份代码包已有647人学习浏览说明其在相关课题中的实操参考价值较高。通过这份代码读者可以直观理解卷积神经网络在CUB-200任务上的完整落地流程从原始图像与标签的解析到卷积网络的结构搭建再到训练与验证的执行各环节均有对应脚本实现整个代码结构清晰适合作为细粒度图像分类的入门基线也可快速修改后迁移到其他鸟类或物体识别数据集为后续实验省去从零搭建的麻烦。整体而言这是一套紧凑而完整的可运行示例尤其适合算法验证与课程设计场景。1. CUB200 是什么十多年没退役的鸟类细粒度基准适合用 5994 张图检验模型的细节分辨力CUB200全称 Caltech-UCSD Birds-200-2011代码里常写成 CUB200 或 CUB_200_2011是细粒度图像分类里最常被拿来当靶子的鸟类数据集200 个物种、11788 张图片官方把其中 5994 张划为训练集、5794 张划为测试集。你拿到的 CUB.zip 可能是官方压缩包也可能是社区或同事二次打包的版本比如带着 doublesi7 这类自己加的标识但内部文件结构万变不离其宗。这个数据集的难点不在于类别多而在于类间差异太小——北美不同种的莺鹀在外观上就差几根羽毛的颜色和一道眉纹模型很容易把注意力放到背景和姿态上所以它被用来验证细粒度分类、迁移学习和注意力机制是否真的有效。适合谁刚跑通 ImageNet 想进一步做细粒度任务的工程师、需要做鸟类生态识别的落地项目以及想检验某个预训练模型表达力的人。2. 把 CUB.zip 解成能训练的状态官方文件、目录结构与标签对齐脚本2.1 解压后先认清官方目录结构别急着写代码先别急着把图片塞进ImageFolder。CUB200 和 ImageNet 最大的不同是它不按一个类一个文件夹 随机划分来组织而是用一组文本文件描述每张图片的归属和划分。解压并查看一级目录unzip CUB.zip -d ./cub200 cd cub200 ls -la正常会看到一个类似CUB_200_2011/的根目录里面有images/图片目录和一批.txt元数据文件。需要认准的是这五个文件内容每行格式images.txt全部 11788 张图片的清单图片编号 相对路径image_class_labels.txt每张图片的类别编号图片编号 类别编号(1–200)train_test_split.txt每张图片的划分标记图片编号 1(训练)/0(测试)classes.txt200 个类别的名称类别编号 拉丁学名bounding_boxes.txt每张图的框坐标图片编号 x y width height注意bounding_boxes.txt里的坐标是相对原图的单位是像素。细粒度分类里常有人拿它裁剪出鸟体区域来减小背景干扰后面第 4 章会提到怎么用。images/里是 200 个以 4 位数字开头的子目录例如001.Black_footed_Albatross。目录名不干净不要直接拿它当类别标签。如果解压后第一眼看到的不是这些文件而是直接一个images目录说明这份 CUB.zip 是有人重新打过包的可能连元数据文件都合并或改名了。我一般会先执行一个自检数一下图片数量和 txt 行数是否都等于 11788。不一致就说明这份包被动过手脚训练前必须修正否则后面标签对不上很容易翻车。2.2 写一个 prepare 脚本把三个 txt 合并成一份干净的清单无论后续用什么框架训练第一步都是把官方文件的图片索引、类别、划分标记对齐成一张表。官方三个 txt 文件的图片编号顺序是一致的所以可以直接按下标合并。在项目根目录建prepare_cub.pyimport os root ./cub200/CUB_200_2011 def read_list(path): items [] with open(path) as f: for line in f: items.append(line.strip().split()) return items images read_list(os.path.join(root, images.txt)) labels read_list(os.path.join(root, image_class_labels.txt)) split read_list(os.path.join(root, train_test_split.txt)) assert len(images) len(labels) len(split) 11788 train_f open(cub_train.txt, w) test_f open(cub_test.txt, w) for img, lab, sp in zip(images, labels, split): img_id, img_path int(img[0]), img[1] cls_id int(lab[1]) # 类别编号 1~200 is_train int(sp[1]) # 1 训练0 测试 abs_path os.path.join(root, images, img_path) line f{img_id}\t{cls_id}\t{abs_path}\n if is_train: train_f.write(line) else: test_f.write(line) train_f.close() test_f.close() print(train:, 5994, test:, 5794)这段逻辑很直白三个列表按下标对齐各自的第二列是图片路径、类别编号、是否训练。把结果写成制表符分隔的文本训练和测试分开放后面 DataLoader 直接读这个文件即可。为什么不用train_test_split.txt之外的方式再随机切分一次因为 CUB200 官方划分是约定俗成的标准几乎所有已发表论文都按这份划分报准确率。你如果自己重切数字就没法跟别人对比了训练集多几张少几张也直接影响收敛行为。2.3 清洗路径与后缀Windows 和二次打包带来的两个坑官方images.txt里的路径用的是 Unix 分隔符在 Linux/macOS 下没问题Windows 下os.path.join有时会混进反斜杠建议统一转成正斜杠再写进清单。另外二次打包的压缩包里图片后缀不一定全是.jpg有的可能是.JPG甚至有人批量转存后图片损坏。写清单时顺手做一步校验from PIL import Image # 抽样检查前 200 张训练图 with open(cub_train.txt) as f: sample [line.split()[2] for line in f.readlines()[:200]] for p in sample: try: Image.open(p).load() except Exception as e: print(broken:, p, e)这一步不耗时但能救命。我拿到一份 CUB.zip 时发现001.Black_footed_Albatross下混了几张.jpeg后缀的图PIL 默认能打开但如果做后端字节校验就会现原形。建议把检查范围放到全部 11788 张几百毫秒而已发现问题直接在清单里替换路径。3. 用 PyTorch 加载 CUB200自定义 Dataset、Transform 与 batch 组织3.1 为什么不能直接 ImageFolder一套标签错位事故还原很多人第一次训练 CUB200 时顺手用torchvision.datasets.ImageFolder(images, transform...)结果训练准确率停在一个很低的值上。原因在于ImageFolder会按子目录名字符串排序生成类别映射而 CUB200 官方classes.txt里的类别编号对应的是生物学分类顺序不是字母序。比如001.Black_footed_Albatross和002.Laysan_Albatross按数字前缀排序没问题但一旦有人把目录重命名成鸟类英文名映射就全乱了。最稳妥的做法是自己写 Dataset以第 2 章生成的cub_train.txt/cub_test.txt为唯一真源。import torch from PIL import Image from torch.utils.data import Dataset class CUBSimple(Dataset): def __init__(self, list_file, transformNone): self.samples [] # (path, label index 0~199) with open(list_file) as f: for line in f: parts line.strip().split(\t) path parts[2] label int(parts[1]) - 1 # 官方标签 1~200 转成 0~199 self.samples.append((path, label)) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label这里最关键的是int(parts[1]) - 1。官方类别编号从 1 开始而 PyTorch 分类头要求 0~199漏掉这一步轻则训练时 loss 不降重则直接越界报错。Image.open(...).convert(RGB)是为了把灰度图和带 alpha 通道的图统一成三通道CUB 里少数老照片是灰度不做这步会在 collate 时因为尺寸不一致直接崩。3.2 细粒度任务的 Transform分辨率、裁剪尺度与归一化CUB200 的官方原图大部分在 300~500 像素之间直接缩到 224 会丢掉羽毛纹理。常见做法是用 448 或 336 作为训练分辨率配合随机裁剪和遮挡增强。下面这组 transform 是我跑 CUB200 常用的起点from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(size448, scale(0.7, 1.0), ratio(0.75, 1.333)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(10), transforms.ColorJitter(0.2, 0.2, 0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) test_transform transforms.Compose([ transforms.Resize(int(448 * 1.15)), # 先放大到 515再中心裁剪 transforms.CenterCrop(448), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])重点参数解释RandomResizedCrop的scale从默认的(0.08, 1.0)收窄到(0.7, 1.0)原因是鸟在画面中一般都占主体过度缩小裁剪会让模型学到从大树背景里猜物种而不是看鸟本身。size参数全链路统一用 448输入分辨率直接决定细粒度准确率上限这一点在 CUB200 上非常明显224 输入下 ResNet50 大概能到 80%448 能到 84% 以上。测试时先Resize到 515 再CenterCrop(448)是为了让鸟体在中心区域占的比例与训练一致也可避免因为原图分辨率不均匀导致缩放后鸟太小。3.3 DataLoader 参数与 batch 组织worker、pin_memory、drop_last数据管道搭好后DataLoader 的参数决定训练速度与 BN 稳定性from torch.utils.data import DataLoader train_loader DataLoader( train_ds, batch_size32, shuffleTrue, num_workers8, pin_memoryTrue, drop_lastTrue, # 丢掉最后一个不完整的 batch稳定 BN persistent_workersTrue ) test_loader DataLoader( test_ds, batch_size32, shuffleFalse, num_workers8, pin_memoryTrue, drop_lastFalse )num_workers在 Linux 上可以开到 8~16Windows 上建议 4 以内并且要把主脚本包进if __name__ __main__:不然会反复 spawn 出多个 dataset 实例。drop_lastTrue只是去掉训练集最后不满一个 batch 的尾巴CUB 训练集 5994 张32 的 batch 会剩 10 张丢掉不影响统计意义却能避免 BN 在最后一个 batch 上统计量抖动。persistent_workers能减少每个 epoch 重建 worker 的开销但显存紧张时要慎重因为每个 worker 会把一张图片的解码缓冲留在内存里。4. 训练 CUB200-2011ResNet 预训练微调的参数设置与训练脚本4.1 模型选型为什么从 ResNet 系列开始而不是直接上 ViTCUB200 单类别训练图只有约 30 张属于典型的小样本细粒度分类完全从随机初始化训练会严重过拟合。业界最常见的做法是取 ImageNet 预训练的 ResNet50 或 ResNet101 作为骨干把最后全连接层换成 200 输出。ResNet 在这个数据集上的表现足够稳定ResNet50 448 分辨率大概在 84%ResNet101 能到 86% 左右这是很多细粒度论文的 baseline 参考线。ViT 也可以做但不建议头一次就上它的预训练权重输入分辨率、patch size 和微调策略都更敏感调不好会一路翻车。import torch.nn as nn from torchvision import models model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V2) in_features model.fc.in_features model.fc nn.Linear(in_features, 200)weights参数用IMAGENET1K_V2这是 torchvision 官方精度更高的一组权重比老默认权重好一点。换掉fc后新层的权重是随机初始化的梯度会比主干大不少所以训练里要么给分类头单独设学习率要么在前几个 epoch 用 warmup 平滑过去。还有个更省事的策略全局统一学习率但把fc的weight_decay调小防止分类头过拟合。4.2 训练策略SGD 还是 AdamW、学习率、epoch 数怎么定CUB200 上我试过两组配置都有稳定效果配置项方案 ASGD 派方案 BAdamW 派优化器SGD(momentum0.9)AdamW初始学习率0.01全层微调1e-4batch size3232输入分辨率448336训练 epoch6030weight decay1e-45e-4学习率调度Cosine 衰减到 0Cosine 衰减到 0label smoothing0.10.1SGD 方案在 CUB200 上收敛更慢但结果普遍更稳尤其是用 ResNet 时微调阶段的主干部分用 SGD 不容易跑飞。AdamW 方案适合只想快速验证数据管道的人30 个 epoch 就能到接近 80% 的水平但上限通常比 SGD 略低。全层微调时初始学习率从 0.01 开始如果发现 loss 震荡先降一半再说只微调最后一两层的情况下主干可以直接不更新但那样效果会差不少。完整训练循环抽出核心部分import torch from torch.optim import SGD, lr_scheduler criterion nn.CrossEntropyLoss(label_smoothing0.1) optimizer SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) scheduler lr_scheduler.CosineAnnealingLR(optimizer, T_max60) for epoch in range(60): model.train() for images, labels in train_loader: images, labels images.cuda(), labels.cuda() outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() scheduler.step() # 每个 epoch 后跑一次验证记录 top1/top5clip_grad_norm_是我在细粒度训练里必留的一行。换掉fc后新头梯度大偶尔一个 outlier 样本会把整个主干权重顶偏clip 到 5.0 基本不影响正常收敛却能避免 loss 突然跳到 20 以上的事故。CosineAnnealingLR的T_max与总 epoch 对齐如果用 AdamW 方案T_max30。4.3 边界框真能用吗一个简单的 Region 裁剪增强不少人在 CUB200 训练时会直接无视官方bounding_boxes.txt其实它是个低成本白嫖的好东西。做法很简单按第 2 章的索引读入每张图对应的框坐标在训练 transform 的RandomResizedCrop之前把图先裁到框范围并外扩 20%再走正常的随机裁剪。这样模型看到的基本都是鸟体主干背景干扰直接少一半。def crop_by_bbox(img, bbox, expand0.2): x, y, w, h bbox cx, cy x w / 2, y h / 2 nw, nh w * (1 expand), h * (1 expand) x1 max(0, int(cx - nw / 2)) y1 max(0, int(cy - nh / 2)) x2 min(img.width, int(cx nw / 2)) y2 min(img.height, int(cy nh / 2)) return img.crop((x1, y1, x2, y2))注意expand不要大于 0.3CUB 的框有的标得比较紧扩太多会把旁边树杈收进来。推理阶段不要用这个裁剪因为测试集框同样存在但真实应用里通常没有框可用训练时用了框、测试时不用会导致分布不一致。最好在测试时也保留整图用CenterCrop让模型自己学会定位鸟体。5. 训练 CUB 的 5 个常见坑从标签错位到 BN 崩溃的排查清单5.1 训练准确率卡在 30% 出头loss 还正常下降现象训练 loss 从 4 左右稳步降到 2但 top-1 一直只有 30% 上下看起来像模型学到了什么却又没全学会。原因十有八九是标签映射错位。前面说过ImageFolder按目录名排序而官方类别编号按生物学分类排列两者对不上。但还有一种更隐蔽的情况你自己写的 Dataset 读image_class_labels.txt时没有-1导致类别 200 的样本标签是 200而你分类头输出维度是 200越界的标签被CrossEntropyLoss静默忽略或随机命中训练集准确率自然会低。解决先打印dataset[0]和dataset[-1]确认标签都落在 0~199 之间再随机抽 10 张图核对路径对应关系。5.2 fine-tune 时改动主干学习率后 BN 统计量直接崩溃现象全层微调跑了十几个 epoch 后验证 loss 突然拉升训练 loss 却还在下降或者 loss 变成 NaN。原因BN 层的统计量在 batch size 较小时比如 8 以下噪声很大主干预训练权重中的 running_mean/var 被小 batch 持续污染一旦某个大梯度能级异常BN 统计量失稳整个网络就崩了。这个在 CUB 上很容易出现因为 448 分辨率撑不起太大 batch。解决一是尽量把 batch size 提到 16 以上二是显存不够就先把主干冻结只训分类头等分类头收敛后再解冻主干三是给 BN 层单独设低学习率或者干脆在解冻前把 BN 层requires_gradFalse。我通常的做法是训练初期跑 5 个 epoch 的 warmup从 1e-4 线性升到目标学习率给 BN 一个缓冲期这算是血泪经验换来的。5.3 验证集用了训练 transform准确率忽高忽低现象同一份权重验证集每次跑出来的准确率差 2~3 个百分点而且不稳定。原因test_transform里误带了RandomResizedCrop或RandomHorizontalFlip验证阶段每跑一次就重新随机一次结果自然忽高忽低。解决验证 transform 只用确定性操作ResizeCenterCrop不要加任何随机项。每个 epoch 验证前用torch.cuda.empty_cache()清一下显存避免验证时显存不足触发回退到 CPU。5.4 提前下载的预训练权重加载报错Key 对不上现象model.load_state_dict(torch.load(resnet50.pth))报 missing key / unexpected key或者数量对不上。原因这份 CUB.zip 的配套训练脚本如果是用timm或旧版 torchvision 保存的权重key 命名和当前版本对不上另外有人把module.前缀留在 state_dict 里。解决加载时strictFalse先看 missing 和 unexpected 分别是什么再写个小脚本把module.前缀剥掉state torch.load(resnet50.pth, map_locationcpu) if any(k.startswith(module.) for k in state.keys()): state {k.replace(module., ): v for k, v in state.items()} model.load_state_dict(state, strictFalse)5.5 混合精度训练下 loss 呈现周期性尖峰现象用torch.cuda.amp训练loss 每隔几百步跳高一次然后又恢复最终准确率比纯 FP32 低 2 个点。原因CUB 的 448 分辨率下小幅值梯度相比 ImageNet 更常见AMP 在缩放 loss 时如果不做动态更新或者模型里有不稳定的分支会出现梯度下溢。细粒度任务的损失曲面本身也窄一点对梯度精度更敏感。解决不直接放弃 AMP而是在GradScaler里显式设定init_scale512.0并把growth_interval从默认 2000 降到 500。还有一个更简单的做法把RandomResizedCrop的scale下限调回 0.5适度降低细节难度AMP 的稳定性会明显改善。6. 验证不只盯着 Top-1 准确率混淆矩阵与 Grad-CAM 把模型糊弄过关的细节揪出来训练结束只看测试集 Top-1 是很危险的。CUB200 有 200 类一个 ResNet50 模型达到 84% 时仍有大约 900 张图被错分。这些错分点才是细粒度模型真正要面对的问题同类不同种、亚种外观重叠、幼鸟与成鸟差异。验证阶段我会做三件事。第一件输出混淆矩阵按测试集预测和真实标签统计重点关注那些误判成关系最近物种的样本对from sklearn.metrics import confusion_matrix import numpy as np all_preds, all_labels [], [] model.eval() with torch.no_grad(): for images, labels in test_loader: images images.cuda() preds model(images).argmax(1).cpu().tolist() all_preds preds all_labels labels.cpu().tolist() cm confusion_matrix(all_labels, all_preds) np.set_printoptions(suppressTrue) # 找出互相关错最多的类别对 pair_max np.unravel_index(cm.argmax(), cm.shape) print(most confused pair:, pair_max)第二件统计每个类别的测试样本中被正确分类的比例把最低的 10 个类别打印出来去images/里翻一翻这些类的原图。我做过一次发现表现最差的是几种翼色和体型都极其接近的柳莺而模型犯错往往统一指向同一类——说明它对关键视觉分叉点的利用不足。第三件用轻量 Grad-CAM 看热力图。做法是取layer4最后一层卷积的梯度求通道平均再上采样原图叠加。不需要引入额外的库torchvision 的feature_maps配合 hook 就能拿到class CAMWrapper(nn.Module): def __init__(self, model, target_layer): super().__init__() self.model model self.fm None self.grad None target_layer.register_forward_hook(self.save_fm) target_layer.register_full_backward_hook(self.save_grad) def save_fm(self, m, i, o): self.fm o.detach() def save_grad(self, m, gi, go): self.grad go[0].detach() def forward(self, x): return self.model(x) def cam(self, class_idx): grad self.grad.mean(dim(2, 3), keepdimTrue) weight self.fm * grad cam weight.sum(dim1, keepdimTrue).relu() return cam.squeeze()如果热力图大面积落在鸟身外比如集中在树枝、水面说明模型在拿背景当线索硬猜。CUB200 的测试集里背景和物种并不完全独立这种模型 Top-1 可能不低但换到真实场景部署时准确率会大幅回落。我自己的习惯是每个实验的最后都跑一遍热力图把问题最大的 5 类记录下来很多次它比准确率更能说明模型学没学到本质。希望这些步骤能帮你在 CUB200 和其他细粒度数据集上少走点弯路。本文还有配套的精品资源点击获取