简介:基于ResNet与CBAM注意力机制实现的Stanford Dogs识别分类Python源码包,面向计算机视觉入门与课设毕设场景,可直接训练并迁移到其他图像分类任务。包内共21个文件,涵盖6个Python脚本、6个编译缓存文件、3个Markdown说明文档、2个shell训练脚本及示例图片,压缩包仅228KB,目录结构清晰,便于快速阅读和二次开发。目前已有383人学习使用。源码实现ResNet50结合CBAM注意力模块的完整训练流程,附带BAM注意力机制对比实现、数据加载与训练启动脚本,并提供介绍文档和更新日志。项目代码经验证可稳定运行,适合作为期末大作业、课程设计或毕业设计初稿,也能替换数据集DIY其他识别功能。借助完整代码和说明,使用者可快速复现Stanford Dogs分类任务,深入理解注意力机制与残差网络的集成方法。
1. 同样是 ResNet50,加上 CBAM 之后在 Stanford Dogs 上为什么还能再涨点
同样是 ResNet50,把 CBAM 注意力模块塞进去之后,Stanford Dogs 上的识别准确率往往能再往上走一到两个点,这不是玄学,而是把“该看什么特征”和“该在哪里看”显式地教给了网络。基于 ResNet + CBAM 注意力机制实现 Stanford Dogs 识别分类的 Python 源码方案,核心是完成 120 个犬种的细粒度分类,同时把数据接口做了通用化处理,换成你自己的图片数据集时,改动主要集中在分类头和数据读取两层。适合用 PyTorch 做图像分类、想在不大改主网络的前提下稳定提点的人群。新手可以直接把训练流程跑通,熟手可以拿它当迁移学习的底座,改成自己的分类版本。整个方案的关键不在于 ResNet 本身,而在于 CBAM 用极小的参数代价换来了可解释的注意力分布。
2. 拆解 CBAM 与 ResNet 的配合:细粒度分类为什么吃这一套
2.1 细粒度分类的难点:狗品种差异比你想的小
Stanford Dogs 的难点和 ImageNet 这种通用分类完全不一样。120 个犬种之间很多差异落在耳朵形态、毛发纹理、尾巴姿态、体态比例这些局部区域上,而不是整体物体轮廓。比如短毛犬种的纹理差异可能只有几条边缘的走向不同,或者脸部的斑点分布密度略有区别。
普通 ResNet 经过四个 stage 的下采样后,特征图分辨率从 224 一路降到 7x7,进入全局平均池化时,空间位置信息已经被压缩得所剩无几。全局平均池化的本质是求特征图的均值,这会把局部判别性区域的强响应淹没在大面积背景或非判别区域里。换句话说,网络在最后一个卷积层“知道”某块区域有响应,但池化时把它和其他区域的响应做了平均,响应强度被拉平了。
细粒度分类场景下模型经常翻车的原因就在这里:相近犬种间的判别信息往往集中在极小的局部区域,而普通 ResNet 的全局池化结构天然不擅长保留这种稀疏强响应。增加网络深度可以缓解,但参数量和显存代价都上去了。相比之下,在特征图上施加注意力权重,让网络在池化之前先显式增强关键区域、抑制无关区域,是一个性价比高得多的做法。
我一般把注意力模块理解为“池化前的一道门控”:它不改变网络的局部感知能力,但改变了哪些信息最终被送去池化。对 Stanford Dogs 这种类别多、类间差异小、类内差异大的数据,这道门控几乎总是能带来可复现的涨点。
2.2 CBAM 的两个子模块:通道注意力与空间注意力的分工
CBAM 全称是 Convolutional Block Attention Module,由通道注意力和空间注意力两部分串联组成。它解决的问题可以拆成两个问题:网络该关注哪种特征?这些特征出现在哪里?通道注意力回答第一个问题,空间注意力回答第二个问题。
通道注意力的计算方式:对输入的 H x W x C 特征图,分别做全局平均池化和全局最大池化,得到两个 1x1xC 的向量。两个向量送入同一个共享 MLP(两层全连接,中间维度压缩为 C/r,r 是缩减率,典型值 16),输出逐元素相加后过 sigmoid,得到一组 0 到 1 的通道权重。平均池化保留特征的全局统计信息,最大池化保留响应最强的局部信号,两者互补,能比只用平均池化更稳定地捕捉判别性通道。
空间注意力的计算方式:把经过通道注意力加权后的特征图沿通道维分别做平均和最大,得到两个 HxWx1 的平面,拼接成一个 2 通道的 HxWx2 图,送入一个 7x7 卷积,降到 1 个通道,过 sigmoid 得到空间权重。7x7 的大卷积核是为了获得较大的感受野,让空间权重能看到周围更广的上下文,避免只聚焦在单像素局部。
两个子模块的顺序是固定的,先通道后空间。设计上的直觉是:先告诉网络“看什么类型的特征”,再告诉网络“这些特征在哪个位置。通道注意力在前,输出是一个通道权重向量,不会引入太多空间模糊;空间注意力在后,利用已经强化过的通道特征去定位关键区域。顺序反过来时,空间注意力会在包含大量噪声通道的原始特征图上计算权重,效果明显下降。
| 模块 | 输入尺度 | 执行方式 | 输出含义 |
|---|---|---|---|
| 通道注意力 | H x W x C | 双池化 + 共享 MLP + Sigmoid | C 维权重,表示各通道重要程度 |
| 空间注意力 | H x W x C | 通道维双池化 + 7x7 卷积 + Sigmoid | H x W 权重,表示各位置重要程度 |
这个模块的参数量非常克制。通道注意力只有一个两层 MLP,空间注意力只有一个 7x7 卷积,两个加起来的参数量相对 ResNet 主干几乎可以忽略。下面第 4 章会把这些模块改成 PyTorch 代码,塞进 ResNet50。
2.3 ResNet 残差块的三种嵌入位置:参数、计算量与效果的取舍
CBAM 可以插在 ResNet 的不同位置上,位置不同效果差别很大。我见过三种常见做法:
第一种是塞进每个残差块内部,通常是残差分支的最后一个 3x3 卷积之后、与 shortcut 相加之前。这种做法的控制粒度最细,每个 block 都做一次注意力加权,效果理论上最强,但问题有两个:浅层特征图分辨率大,注意力计算量成倍上升;浅层特征本身语义信息弱,在早期就做通道加权容易把低级纹理特征过早地筛选掉,后期想恢复都难。
第二种是接在每个残差块输出之后、ReLU 激活之前。输出经过注意力加权后再进入激活,注意力相当于门的角色。这种做法比内部插入更简洁,但每个 block 后都挂一个 CBAM,参数量和显存开销同样不小,对 ResNet50 全量插入时,训练显存和推理时延都会有明显上浮,收益相对不那么划算。
第三种,也是我在这套方案里采用的做法:只在高层 stage 之后挂 CBAM,也就是 ResNet50 的 layer3 和 layer4 输出之后。理由是高层特征已经具备语义信息,狗的耳朵、毛色、体态这些判别性概念在这里已经成型,在全局池化之前做最后一次注意力加权,效率和收益的平衡最好。ResNet50 的 layer3 输出通道是 1024,layer4 输出通道是 2048,两个位置各挂一个 CBAM,参数增量只有百万级左右,相对 ResNet50 两千多万的总参数可以忽略。
选择位置时还要考虑一个实际问题:CBAM 插得越深,梯度回传路径越短,训练越稳定;插得越浅,梯度需要穿过更多注意力计算层,如果主干的预训练权重被冻结,浅层注意力很容易在训练初期产生不稳定的梯度信号。实践里我建议优先复现“layer3 + layer4 末端两个 CBAM”这个配置,跑通之后再尝试逐步加深,而不是一上来就把每个残差块都改造掉,那样出了问题都分不清是注意力的问题还是主干的问题。
3. 准备 Stanford Dogs 数据并规划迁移路径:目录结构决定后续手顺不顺
3.1 下载解压与目录组织:先让训练脚本认识你的数据
Stanford Dogs 的原始压缩包解压后,Images 目录下每个子目录代表一个犬种,目录名是 ImageNet 风格的 wnid 编号加品种名,类似 n02085620-Chihuahua 这种格式。图片数量每个品种不均衡,有的品种一百多张,有的只有六七十张。正式训练前要先把这些文件重新组织成训练集和验证集。
常见的数据加载做法是依赖目录结构自动生成标签,目录名就是类别名。我一般把数据整理成下面这样,train 和 val 下各有 120 个子目录,每个子目录内是该品种的图片,两层目录结构一目了然,也不容易出错。
data/stanford_dogs/ ├── train/ │ ├── n02085620/ │ │ ├── n02085620_10074.jpg │ │ └── ... │ ├── n02099712/ │ └── ... └── val/ ├── n02085620/ └── ...从原始 Images 目录整理到上述结构,我一般用下面的脚本,按 8:2 比例随机划分每个类别,同一个类别的图片先打乱再切分,避免某个品种整体全落在验证集里:
import os import shutil import random from pathlib import Path source = Path("Images") train_root = Path("data/stanford_dogs/train") val_root = Path("data/stanford_dogs/val") val_ratio = 0.2 random.seed(42) for cls_dir in source.iterdir(): if not cls_dir.is_dir(): continue imgs = list(cls_dir.glob("*.jpg")) random.shuffle(imgs) split = int(len(imgs) * (1 - val_ratio)) for img in imgs[:split]: dst = train_root / cls_dir.name / img.name dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy(img, dst) for img in imgs[split:]: dst = val_root / cls_dir.name / img.name dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy(img, dst)这段脚本的逻辑是:遍历 Images 下每个品种目录,读取全部 jpg 文件,打乱后按比例切出验证集,然后复制到目标目录。用shutil.copy而不是os.rename,是因为我不想破坏原始下载数据,复制出错还可以重新来过。val_ratio = 0.2表示每个类别都要切 20% 的图片给验证集,这样验证集的类别分布和训练集一致,不会出现某个品种在验证集里没有样本的情况。random.seed(42)是固定随机种子,保证任何两次运行得到相同的划分结果,方便复现和后续排查。
这里有个容易出错的细节:脚本里用的是cls_dir.name作为类别目录名,原始目录名包含 wnid 编号和品种名,中间有个连字符。ImageFolder加载时会把整个目录名当作类别名,所以class_to_idx映射中键值会是完整的n02085620-Chihuahua这种格式。推理时如果要显示中文品种名,需要额外加一个映射字典,把 wnid 对应到日常名称。
3.2 换成自己的数据集:最少改动量的四个改造点
标题里特别强调了“可换其他数据识别分类”,这是这套方案最有价值的部分。实际换数据时不需要改模型结构,只需要改四个地方。
第一是图片目录结构。把自己的数据按照和上一节相同的两层目录组织好,即train/类别名/图片、val/类别名/图片,类别的目录名用什么都可以,最终都会成为模型的标签名。如果原始数据是 CSV 格式的标注,建议先写一个脚本把图片按标签移动到对应目录下,这样后面所有代码都不用为数据格式做适配。
第二是类别数num_classes。模型最后一层全连接输出维度必须等于自己的类别数。Stanford Dogs 是 120,换自己的数据时改成实际类别数量。最常见的问题是想当然地把类别数写成目录个数,结果某个目录是空的,数量对不上,加载模型时直接报维度不匹配。
第三是图像通道和归一化参数。Stanford Dogs 的图都是 RGB 三通道,直接沿用 ImageNet 预训练模型的归一化参数即可,也就是红绿蓝三通道各自的均值和标准差。如果你的数据是灰度图,要么把单通道复制成三通道输入,要么改写模型第一层卷积的输入维度;如果你的图像内容本身和自然图像差异极大,比如医学影像或工业检测图像,沿用 ImageNet 的归一化统计量可能不是最优的,建议在自己的数据上重新统计均值和标准差。
第四是数据增广策略。Stanford Dogs 两万多张图对 120 个类别来说并不充裕,训练时一般会做随机裁剪、水平翻转、轻微色彩抖动这些增强;换成自己的数据时,图片尺寸和类别数量变了,裁剪尺度可能要重新调,比如你的图片主体很小,就不能做中心大比例裁剪,否则主体直接被切掉。
这四个改造点的对比如下:
| 改造点 | Stanford Dogs 原值 | 换成自己的数据集 |
|---|---|---|
| 图片目录 | Images/按品种分目录 | train/类别/图片两层结构 |
| 类别数 | 120 | 实际目录数量 |
| 归一化参数 | 自然图像三通道统计量 | 灰度图或特殊分布需重算 |
| 输入尺寸 | 224x224 | 按数据实际主体占比调整 |
换数据之后,最稳的验证路径是:固定随机种子,先跑通 5 个 epoch,确认 loss 能稳定下降、验证集准确率超过随机猜测的100/类别数这条基线,再把它当成正式训练前的可用性检查。
3.3 训练/验证划分与类别映射:别让标签对不上
用 PyTorch 的ImageFolder加载数据时,类别索引是按目录名的字母序自动生成的。训练集和验证集目录如果来源相同,生成的class_to_idx通常一致,但不能默认一致,最好在训练前显式断言一次,避免隐藏的对齐问题:
from torchvision import datasets train_set = datasets.ImageFolder("data/stanford_dogs/train") val_set = datasets.ImageFolder("data/stanford_dogs/val") assert train_set.classes == val_set.classes, "train/val 类别不一致" assert len(train_set.classes) == 120, "类别数不是 120" num_classes = len(train_set.classes) print(train_set.class_to_idx)这段脚本的作用是加载两个数据集的目录结构并建立标签映射,然后断言两边类别列表完全一致。这个检查非常关键:如果训练集 120 类而验证集漏了一个类,ImageFolder不会报错,但模型输出的 120 维 logits 和验证集的标签编号会错位,验证准确率会莫名其妙地偏低,又很难排查。assert len(train_set.classes) == 120是把类别数作为一个硬性约束写死,一旦有人改乱了数据目录就会第一时间暴露。
ImageFolder的索引用于CrossEntropyLoss时没有额外问题,但要注意它生成的标签是 0 到 119 的整数,和模型输出的类别索引是一致的。如果自定义了标签读取逻辑,比如从 CSV 文件读标签,务必确认 CSV 中的类别名和目录名完全匹配,大小写或空格不同都会导致标签错位,这是换数据时最隐蔽的坑。
4. 用 PyTorch 把 CBAM 嵌入 ResNet:以 ResNet50 为例跑通训练
4.1 CBAM 模块的 PyTorch 实现:通道抑制与空间加权
CBAM 的实现不复杂,完整版本不到 60 行。通道注意力用双池化加共享 MLP,空间注意力用拼接加 7x7 卷积,两个子模块的代码可以独立看,先实现通道注意力:
import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, in_planes, reduction=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.shared_mlp = nn.Sequential( nn.Conv2d(in_planes, in_planes // reduction, 1, bias=False), nn.ReLU(inplace=True), nn.Conv2d(in_planes // reduction, in_planes, 1, bias=False), ) def forward(self, x): avg_out = self.shared_mlp(self.avg_pool(x)) max_out = self.shared_mlp(self.max_pool(x)) return torch.sigmoid(avg_out + max_out)这段代码里用两个自适应池化把任意尺寸的特征图压到 1x1,池化后的特征图经过共享 MLP 得到通道权重。特别注意shared_mlp用的是 1x1 卷积而不是nn.Linear,这样输入输出都保持四维张量的形状,省去手动squeeze和unsqueeze的麻烦。reduction=16表示中间通道数压缩为原来的十六分之一,ResNet50 的 layer3 输出是 1024 通道,中间层就是 64 通道,参数量很小。bias=False是因为后面接 ReLU 和输出层,去掉偏置能减少过拟合风险,也是这个模块的常见配置。
空间注意力模块的输入不是原始特征图,而是经过通道注意力加权后的特征图:
class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False) def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out = torch.max(x, dim=1, keepdim=True).values return torch.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))空间注意力的逻辑是:沿通道维分别取平均和最大值,得到两个单通道平面,拼成 2 通道特征图后送入 7x7 卷积。padding=kernel_size // 2保证输出尺寸和输入一致,这里 7x7 卷积对应padding=3。torch.max(x, dim=1, keepdim=True)返回的是最大值张量,.values取出结果,注意不要用torch.max的返回元组直接参与卷积,容易出类型错误。
把两个模块串起来就是完整的 CBAM:
class CBAM(nn.Module): def __init__(self, in_planes, reduction=16, kernel_size=7): super().__init__() self.ca = ChannelAttention(in_planes, reduction) self.sa = SpatialAttention(kernel_size) def forward(self, x): x = self.ca(x) * x x = self.sa(x) * x return x这个模块的 forward 逻辑很直白:先乘通道权重,再乘空间权重。注意reduction的值决定了 MLP 中间层宽度,如果in_planes // reduction小于等于 0 就会报错,所以当输入通道数很小时要调小 reduction,例如通道数是 64 时,reduction 用 8 而不是 16。
4.2 把 CBAM 塞进 ResNet50:用 Hook 实现不改源码的最小侵入
我已经在前面解释过选择“layer3 末端 + layer4 末端”这个位置。实现方式有两种:修改Bottleneck类源码,在特定阶段插入 CBAM;或者用register_forward_hook,在不改任何 ResNet 源码的前提下把 CBAM 挂在层输出之后。hook 方式更符合工程习惯,后续想换 ResNet18 或 ResNet101 时,代码不需要改动,只需把主干模型换掉即可:
import torchvision.models as models from torchvision.models import ResNet50_Weights model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1) cbam_layer3 = CBAM(in_planes=1024) cbam_layer4 = CBAM(in_planes=2048) def make_hook(cbam): def hook(module, inputs, output): return cbam(output) return hook model.layer3.register_forward_hook(make_hook(cbam_layer3)) model.layer4.register_forward_hook(make_hook(cbam_layer4)) model.fc = nn.Linear(2048, 120)这段代码的逻辑是:先加载带 ImageNet 预训练权重的 ResNet50,然后分别在layer3和layer4的输出上挂 forward hook。hook 函数的返回值会替换模块的原始输出,也就是说 ResNet 前向传播时,layer3计算完先经过 CBAM 加权,再进入layer4;layer4计算完经过 CBAM 加权,最后进入全局池化和全连接层。model.fc = nn.Linear(2048, 120)把原来 1000 维的输出改成 120 维,对应 Stanford Dogs 的类别数。
hook 写法上有一个细节:make_hook用参数把cbam绑定到局部变量,避免在循环中注册 hook 时闭包变量被重复使用。这里的cbam_layer3输入通道是 1024,对应layer3的输出通道数;cbam_layer4输入通道是 2048。换用其他 ResNet 版本时,例如 ResNet18 的layer3是 256 通道、layer4是 512 通道,必须同步修改in_planes,否则卷积权重维度对不上会直接报错。
需要注意一个容易疏忽的点:weights=ResNet50_Weights.IMAGENET1K_V1会在运行时联网下载预训练权重。如果运行环境的网络受限,建议提前单独下载权重文件,然后用下面的方式加载:
model = models.resnet50(weights=None) state = torch.load("resnet50_weights.pth", map_location="cpu") model.load_state_dict(state, strict=False)strict=False很关键,因为后面替换了fc层,原权重里fc.weight和fc.bias的维度是 1000 类,和现在模型的 120 类不匹配,strict 模式下加载会直接抛异常,strict=False会忽略这些维度不匹配的键。
4.3 训练脚本的关键参数:优化器、学习率、冻结策略与损失函数
模型搭好之后,训练参数直接影响能不能收敛。Stanford Dogs 只有两万张图,对 120 分类来说数据量偏少,直接全量训练会破坏 ImageNet 预训练模型的特征表达能力。常见做法是冻结部分底层参数,只微调高层。我的配置如下:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) for name, param in model.named_parameters(): if "layer1" in name or "layer2" in name: param.requires_grad = False criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD( [p for p in model.parameters() if p.requires_grad], lr=0.01, momentum=0.9, weight_decay=1e-4, nesterov=True, ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=60)这段配置里lr=0.01是配合冻结浅层使用的。全量微调时建议降到0.001起步,因为所有层都会更新,学习率过高容易把预训练权重一把冲毁。冻结layer1和layer2的理由是这两个阶段提取的是边缘、颜色、纹理这类通用特征,换到狗品种分类时依然适用,不需要重新学习;layer3和layer4以及新增的 CBAM 模块则保留requires_grad=True,让它们去适应细粒度判别。
weight_decay=1e-4是 ImageNet 训练中常用的 L2 正则强度,防止全连接层过拟合。nesterov=True是在标准 SGD 动量基础上使用 Nesterov 加速梯度,收敛更平稳。CosineAnnealingLR的学习率从 0.01 按余弦曲线降到接近 0,整个周期 60 个 epoch,相比 StepLR 的阶梯式下降,这种方法在微调场景下通常更稳,后期不容易在最优解附近反复震荡。
训练循环里还有一个与冻结策略配套的隐性要求:BN 层的运行统计量更新策略。PyTorch 的model.train()会让所有 BatchNorm 层继续更新 running_mean 和 running_var。如果冻结了layer1和layer2,同时被冻结的 BN 层不应当再更新统计量,正确做法是在每次前向传播前,把冻结阶段的 BN 层切到 eval 模式:
for name, module in model.named_modules(): if "layer1" in name or "layer2" in name: if isinstance(module, nn.BatchNorm2d): module.eval()这段设置的逻辑是:冻结参数后,layer1和layer2的卷积核已经不再变化,如果它们的 BN 层还在用当前 batch 的统计量做归一化,那运行统计量每次都在变,而卷积权重不变,两者就不匹配了。切到 eval 模式后,BN 使用冻结的历史统计量,行为一致且稳定。没有这一步,我见过不少训练初期 loss 能降、几个 epoch 后开始剧烈波动的症状,根源都在这里。
训练过程中建议关注两个指标:训练集 top-1 准确率和验证集 top-5 准确率。Stanford Dogs 类别间相似度高,top-5 更能反映模型是否把相近品种排到了前面。验证阶段用torch.no_grad()包住前向计算,同时把模型整体切到model.eval(),但要注意model.eval()会把所有 BN 层都切掉,这会同时影响被冻结的 BN 层,没有副作用,因为验证阶段本来就该用历史统计量。
5. 训练与迁移常见问题排查:5 个最容易翻车的坑
5.1 加载预训练权重报尺寸不匹配
现象:运行脚本时抛异常,提示fc.weight和fc.bias的尺寸不匹配,比如size mismatch for fc.weight: copying a param with shape torch.Size([1000, 2048]) from checkpoint, the shape in current model is torch.Size([120, 2048])。
原因:加载预训练权重时用了默认的strict=True,但模型的fc已经被替换成了 120 类输出,权重文件中 1000 类的参数无法对应。
解决:加载时显式传入strict=False,跳过fc层的权重复制。fc层是随机初始化的,第一次训练时它输出的是随机 logits,损失偏大是正常现象,通常几个 epoch 后全连接层会逐渐收敛,准确率回到正常水平。如果想让fc层初始化更稳一点,可以手动设置:
nn.init.kaiming_normal_(model.fc.weight, mode="fan_out", nonlinearity="relu") nn.init.constant_(model.fc.bias, 0.0)也可以把新的fc层初始化为小数值,但实际经验表明,迁移学习中随机初始化的分类头配合预训练主干,用常规 LR 微调就能很好收敛,过度设计初始化反而多余。
5.2 loss 不降或下降极慢
现象:训练了十几个 epoch,训练集 loss 几乎不下降,验证集准确率停留在随机猜测水平上下。
原因:最常见的是学习率设置不合适。冻结浅层后如果还用lr=0.001这种偏小的值,新增的 CBAM 模块和fc层因为是从随机初始化开始学习的,梯度更新幅度不足。另一个原因是优化器只传入了p.requires_grad的参数,如果冻结策略判断条件写错了,比如把layer3也写进冻结范围,等于只训练随机初始化的fc层,loss 自然降不下去。
解决:先把打印出来的可训练参数量确认一遍。在优化器构造前加一行:
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print("trainable params:", trainable_params)ResNet50 总参数量两千多万,冻结layer1和layer2后,可训练参数量应该还剩接近两千万级别。如果打印结果只有几百万,说明冻结范围过大。学习率方面,冻结浅层时用0.01起步,全量微调用0.001,新增的 CBAM 和分类头如果单独设置学习率倍率,可以给它们乘以 10 的系数,让新模块学得更快一些。
5.3 训练时显存溢出
现象:训练到第一个 forward 或第二个 iteration 时直接报 CUDA out of memory,提示显存不足,有时还会附带Tried to allocate 256.00 MiB这类信息。
原因:CBAM 的注意力计算是逐层保留中间激活的,每经过一次 CBAM,反向传播都需要保存它的输入和中间结果;如果在每个 block 后都挂 CBAM,显存激活值会比原始 ResNet50 高出 30% 以上。这个问题在 batch size 较大或输入分辨率较高时尤其明显。
解决:优先减小 batch size,比如从 64 降到 32,通常能立刻解决问题。第二个手段是梯度累积,用多个小 batch 更新一次参数,模拟大 batch 效果,显存占用不变大:
accum_steps = 4 for i, (images, labels) in enumerate(train_loader): loss = criterion(model(images.to(device)), labels.to(device)) loss = loss / accum_steps loss.backward() if (i + 1) % accum_steps == 0: optimizer.step() optimizer.zero_grad()这段代码把一次参数更新拆成 4 次前反向,loss 除以accum_steps后再累积梯度。注意 BN 层的统计量在每个小 batch 上都会更新,所以梯度累积只等价于“参数更新”层面的大 batch,不代表真正的 batch size 变大,BN 行为依然是小 batch 的。如果冻结了 BN 层,这个问题就可以忽略。第三个手段是检查是否真的需要两个 CBAM,stage 末端挂一个也能有大部分效果,显存会更宽松。
5.4 验证集准确率震荡、训练提前过拟合
现象:训练集准确率稳步上升,但验证集准确率到第 20 个 epoch 后开始剧烈震荡,甚至出现下降趋势,训练集准确率保持高位,差距明显拉大。
原因:Stanford Dogs 每个类别只有一百多张图,模型又带预训练权重,高表达能力容易让分类头在训练后期把训练样本的特殊背景或姿态记进去。加上 120 类输出维度大,分类头参数量多,过拟合风险天然偏高。
解决:数据增强手段分别是随机裁剪、随机水平翻转和色彩抖动。Stanford Dogs 的图片包含大量背景干扰,随机裁剪可以让模型不只依赖背景区域。另一个有效手段是 Label Smoothing,把交叉熵损失中 one-hot 标签放松成平滑分布,防止模型对训练集样本输出过高置信度:
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)label_smoothing=0.1在 PyTorch 中只需要在CrossEntropyLoss里加一个参数即可,不需要改其他任何逻辑。损失函数不再追求 logits 对训练样本的绝对高置信,而保留一小部分概率给其他类别,这能显著缓解验证集上的震荡。配合 CosineAnnealingLR,后期学习率降低后验证准确率通常会回升。
5.5 迁移到自己的数据后某几个类别完全分不出来
现象:换用自己的数据集训练后,整体准确率可观,但查看每个类别的召回率时发现有一两类召回率接近 0,模型基本不认识它们。
原因:最直接的原因是样本数量极不平衡,多数类有几千张,少数类只有几十张,交叉熵损失对多数类有天然偏好。另一个原因是训练集和验证集的划分方式不当,比如按文件名首字母排序后直接切分,导致同一个类别的图片集中出现在某一个集合里。
解决:先修划分逻辑,确保按类别分层抽样,我习惯用sklearn.model_selection.train_test_split配合stratify参数。对于类别不平衡本身,最直接的方法是给损失函数加权:
from sklearn.utils.class_weight import compute_class_weight class_weights = compute_class_weight( class_weight="balanced", classes=train_set.classes, y=[s[1] for s in train_set.samples], ) class_weights = torch.tensor(class_weights, dtype=torch.float32, device=device) criterion = nn.CrossEntropyLoss(weight=class_weights)compute_class_weight会自动计算每个类别的权重,样本少的类别权重更高,样本多的类别权重更低。criterion在损失计算时对少数类的错误加大惩罚,模型会更努力地学到少数类的特征。如果加权后少数类仍然不行,再考虑针对少数类做额外的数据增强,比如复制粘贴、随机旋转角度更大一些,而不是无脑采集新数据。
6. 进阶验证:用 Grad-CAM 检查注意力到底看对没有
训练跑的多了,你会发现准确率只是数字,真正让人安心的是可视化——不看图,你永远不知道模型是靠狗头的耳朵分辨品种的,还是靠背景里的草地蒙对的。Grad-CAM 是这类验证里最成熟的手段,它利用目标类别的梯度对最后一个卷积层的特征图做加权平均,生成一张和输入图像同尺寸的热力图,高亮区域就是模型做出判断的主要依据。
做法分三步。第一步在layer4的输出上注册 forward hook 和 backward hook,分别捕获前向特征图和反向梯度;第二步对梯度取通道维平均,得到每个通道对目标类别的贡献权重;第三步把特征图按权重加权求和,过 ReLU 后归一化到 0 到 1 范围,叠加回原图。核心代码示意如下:
features, gradients = {}, {} def forward_hook(module, inputs, output): features["value"] = output.detach() def backward_hook(module, grad_input, grad_output): gradients["value"] = grad_output[0].detach() model.layer4.register_forward_hook(forward_hook) model.layer4.register_backward_hook(backward_hook) output = model(image.unsqueeze(0)) target = output.argmax(dim=1) model.zero_grad() output[0, target].backward() weights = gradients["value"].mean(dim=(2, 3), keepdim=True) cam = (weights * features["value"]).sum(dim=1, keepdim=True) cam = torch.relu(cam) cam = (cam - cam.min()) / (cam.max() - cam.min())运行这份代码时要注意两个细节:backward hook 拿到的grad_output是梯度张量的元组,取第一个元素参与计算;热力图归一化前一定要先过 ReLU,只保留正向贡献区域,否则会有一些负权重的噪声区域混进来。cam的尺寸和layer4输出一致,即 7x7,直接叠加在 224x224 原图上会显得很粗,我一般先用双线性插值上采样到原图尺寸再叠加。
验证 CBAM 是否起作用的方法很直接:跑一个没有 CBAM 的普通 ResNet50 作为对照,分别导出同一批验证集图片的热力图,对比狗的头部区域是否更集中。CBAM 加了之后,热力图通常会从散落在全身的多块区域收敛到耳朵、眼睛、口鼻这些判别力最强的局部。如果对照实验里两个模型的热力图几乎没有差别,说明注意力模块没有学到新的空间筛选能力,问题大概率出在训练配置上,比如冻结了 CBAM 的参数或者学习率太低。
把这个可视化过一遍再部署模型,比我只看准确率曲线踏实得多。我自己每次跑完这套流程,都会随机抽几十张图翻热力图,确认没有出现“看背景识狗”的偷懒特征,再把它交给下一个环节的人。希望帮到你。
本文还有配套的精品资源,点击获取