2026/9/23 22:09:26

CrossFormer图像分类实战:跨尺度注意力机制与训练避坑指南

CrossFormer图像分类实战:跨尺度注意力机制与训练避坑指南 简介CrossFormer实战资源包面向图像分类方向的开发者与研究人群聚焦跨尺度注意力机制对多尺度特征交互的改进可用于复现分类实验、替换骨干网络或在此基础上改造模型以适应自定义任务。压缩包共2000个文件大小约835.34MB主要包括1986张PNG图像样本、Python训练与推理脚本、pyc缓存文件、类别配置JSON、说明文本以及训练好的pth模型权重构成一套结构完整、目录清晰、可直接运行的项目。已有191人学习下载可作为快速体验CrossFormer图像分类效果的入门物料也适合用于课堂演示或论文对比实验。借助其中的源码与预训练权重读者既可以直接对图片进行分类测试也可以结合配置文件和数据集进行迁移训练或调参研究从而深入理解跨尺度注意力机制的实际应用。整体资源对初次接触视觉Transformer的开发者尤其友好能在较短时间内完成环境验证与结果复现。1. CrossFormer 图像分类这个双流基座模型到底强在哪做视觉 Transformer 分类任务的人大多经历过一种尴尬换了更强的 ViT 变体精度涨了但训练更不稳了尤其碰到多尺度目标混在一起的图片小目标老是丢。CrossFormer 这类双流基座模型核心思路是把多尺度交互直接做进注意力机制里而不是靠堆特征金字塔去补救。它的跨尺度注意力Cross-scale Attention让不同分辨率的 patch 在进入 Transformer 之前就先完成信息交换这也是它和 Swin、PVT 最大的分水岭。这篇笔记我会把 CrossFormer 做图像分类的完整流程拆开讲从跨尺度 patch embedding 的原理、数据预处理、训练脚本到三个必须避开的坑末尾附上注意力可视化的验证手段帮你判断模型是真的学到了特征还是光记住了训练集。2. 跨尺度注意力机制CVA 与 LSA / SDA 的配合逻辑与初始化2.1 CVA 跨尺度 patch embedding小 patch 和大 patch 的融合方式CrossFormer 的第一个核心组件是 CVACross-scale Patch Embedding。传统 ViT 是把图像切割成固定大小的 patch然后做线性映射不同尺度之间没有任何交流。Swin 虽然做了层级化但每个 stage 内部的 patch 尺度是固定的跨尺度交互靠的是下采样和窗口偏移属于间接交互。CrossFormer 的做法更直接每个 stage 同时使用大 patch 和小 patch让它们共享一组线性层输出再拼接起来。从实现角度看CVA 的做法是把输入特征图分别切成小 patch比如 4x4和大 patch比如 8x8 或 16x16小 patch 经过一个卷积层通常是 kernel 为 4、stride 为 4 的 Conv2d大 patch 经过另一个卷积层kernel 为 8、stride 为 8或者同一个卷积层当尺寸存在倍数关系时。然后把两者的输出在通道维度上拼接得到融合后的 token 序列。这里的参数选择直接影响精度和计算量我用一张表格说明常见配置配置项小 patch 分支大 patch 分支输出 token 长度kernel / stride4 / 48 / 8小分支的 4 倍kernel / stride4 / 416 / 16小分支的 16 倍推荐场景小目标密集大目标为主平衡型取 4 / 8实际训练时大 patch 分支不能选得太大否则小目标的细节在 embedding 阶段就被抹掉了。我一般会把图像短边缩放到 224patch 组合取 4 8通道数 64 起步这样在 ImageNet 这种复杂数据上效果比较稳。如果只做单类别或者小数据集比如你的 class.json 里只有 10 类可以大胆取 4 16精度上限更高但需要更强的数据增强来压过拟合。2.2 LSA 局部注意力与 SDA 子采样注意力H 参数和 L 参数的分工CrossFormer 的 Transformer block 内部是双流设计。一条流是 LSALocally-grouped Self-Attention把 token 按窗口分组窗口尺寸由 H 参数控制组内做标准自注意力另一条流是 SDASubsampled Attention先把整个特征图下采样subsample成更小的分辨率然后做全局自注意力。这两条流的输出在通道维度拼接再经过 FFN。这里有一个关键参数H 控制 LSA 的窗口大小L 控制 SDA 的下采样倍数。H 越大局部感受野越大但计算量按 H 的平方增长L 越大SDA 的分辨率越低全局信息的粒度越粗糙。我调试下来H7、L4 是默认平衡点但如果你的输入分辨率不是 224 而是 384H7 会导致窗口边界不齐需要调整 H 使得图像尺寸能被 H 整除。384 的话 H8 更合适。下面是初始化实际模型时的 config 写法用代码块说明 H 和 L 怎么逐 stage 配置import torch from crossformer import CrossFormer # 每层 stage 的深度、H 窗口大小、L 下采样倍数 # 对照原论文的 CrossFormer-S: stages[2,2,6,2] config { in_chans: 3, num_classes: 10, # 你的 class.json 里有多少类就填多少 depths: [2, 2, 6, 2], # 4 个 stage每个 stage 的 transformer block 数 dims: [64, 128, 256, 512], # 每个 stage 的通道数 group_size: [7, 7, 7, 7], # LSA 的 H 参数每个 stage 可不同 subsample_factor: [4, 4, 4, 4], # SDA 的 L 参数 mlp_ratio: 4.0, qkv_bias: True, } model CrossFormer(**config) print(model) # 用一个随机输入测试前向输出尺寸 dummy_input torch.randn(2, 3, 224, 224) out model(dummy_input) print(输出形状:, out.shape) # 期望 [2, 10]如果类别数对得上代码逻辑说明CrossFormer 的构造函数接收group_size和subsample_factor两个列表长度必须等于 stage 数。每个 stage 内部交替使用 LSA 和 SDALSA 的分组大小由group_size决定SDA 的下采样倍数由subsample_factor决定。mlp_ratio控制 FFN 的隐藏层倍数一般用 4.0。qkv_bias建议保持 True尤其是小数据集bias 能提供额外的拟合能力降低收敛难度。参数微调建议如果你的数据集目标尺寸跨度大比如同时有 30x30 的小物体和 300x300 的大物体把 stage 1 的group_size调小到 3 或 5让前面的层更关注局部细节subsample_factor在前两个 stage 保持 4后两个 stage 可以调到 2让深层保留更多空间信息。这个改动在 CIFAR 级别的数据上能提升 1 到 2 个点但训练时间会增加约 15%需要自己权衡。2.3 相对位置编码CrossFormer 的 RPE 是怎么嵌进去的CrossFormer 与 Swin 的另一个差异是位置编码方式。Swin 用的是相对位置偏置表relative position bias table在注意力计算时查表加上偏置值。CrossFormer 用的是可学习的相对位置编码RPE编码是直接在注意力 logits 上叠加的公式上等价于Attention(Q,K,V) softmax(QK^T / sqrt(d) RPE) V这里的 RPE 是一个可学习参数矩阵形状和 attention logits 完全一致。但问题在于LSA 和 SDA 的注意力矩阵尺寸不同LSA 是窗口内的局部注意力形状是[B * num_windows, H, H, C]SDA 是下采样后的全局注意力形状是[B, num_heads, H*W/L^2, H*W/L^2]。所以 RPE 也要分成两套分别对应 LSA 和 SDA。如果你用的是我上面这个crossformer包RPE 的初始化是自动的不需要手动处理。但如果你自己复刻模型或者从别的仓库移植权重最容易踩的坑是 RPE 尺寸不匹配。代码里初始化模型的参数是rpeTrue或rpeFalse注意当输入分辨率改变时RPE 表不会自动插值需要手动调整input_resolution。下面这段代码是处理 224 分辨率图片后再换到 384 分辨率时重置 RPE 的常用做法def reset_rpe_for_resolution(model, new_resolution384): # 遍历所有 block重新生成相对位置编码表 for blk in model.blocks: if hasattr(blk, lsa): # 重新计算 LSA 的 RPE 表根据窗口大小 H blk.lsa.rpe_table torch.nn.Parameter( torch.zeros(blk.lsa.num_heads, blk.lsa.H * blk.lsa.H, blk.lsa.H * blk.lsa.H) ) if hasattr(blk, sda): # 重新计算 SDA 的 RPE 表根据下采样后的分辨率 h w new_resolution // blk.sda.subsample_factor blk.sda.rpe_table torch.nn.Parameter( torch.zeros(blk.sda.num_heads, h * w, h * w) ) return model代码逻辑说明RPE 表和输入分辨率严格绑定换分辨率不重置的话模型会直接报 shape mismatch 错误或者在推理时因为索引越界而静默出错。上面的代码是遍历每个 Transformer block找到 LSA 和 SDA 子模块之后用目标分辨率重新生成零初始化 RPE 表。注意零初始化之后要重新训练几个 epoch 让 RPE 表适应新分辨率不能直接拿去推理。3. 数据集准备与预处理class.json 和图片目录的实战配合3.1 目录结构torchvision 的 ImageFolder 怎么和 class.json 对齐CrossFormer 训练图像分类任务的常规路径是使用torchvision.datasets.ImageFolder来加载图片。它要求数据集按类别放在子目录里子目录名可以是类别名或类别编号。class.json 在这里起的作用是映射「子目录名 - 类别索引 - 类别标签」。举个例子你的数据目录可能是dataset/ train/ 0/ 5e4d1ee0d.png 77291b3ad.png 1/ 0367e0199.png val/ 0/ 5a8b75712.png 1/ 8029e3396.pngclass.json 的内容则可能长这样{ 0: cat, 1: dog, 2: bird }当你用ImageFolder读取目录时它会根据子目录的字典序自动分配索引。有个常见坑如果你的子目录名是字符串标签比如cat、dog排序之后索引是{bird: 0, cat: 1, dog: 2}这和 class.json 里手动写的映射可能不一致。所以我在实战中会先用os.listdir扫一遍目录显式生成 class.json再交给ImageFolder使用。否则模型训练完了推理时类别标签对不上整个流程就是白做。3.2 数据增强策略图像分类用 RandomResizedCrop 还是 Resize CenterCropCrossFormer 对输入分辨率比较敏感因为它有窗口注意力和下采样注意力两者的计算复杂度都和分辨率强相关。我做过一轮对照实验用 224x224 输入、RandomResizedCrop(224, scale(0.08, 1.0))作为训练增强验证集用Resize(256) CenterCrop(224)。这个组合在 ImageNet 上基本是标配但换到自定义数据集时需要关注 scale 的下界。如果你的数据集中小目标多scale(0.08, 1.0)会把目标缩得很小可能导致小目标在 CVA 阶段就丢失。这时把下界调高到 0.3 更安全。另外 CrossFormer 的双流注意力在训练时需要更宽的增强范围来防止过拟合我建议加上RandAugment(n2, m9)和RandomErasing(p0.25)。RandAugment 的m值不要超过 10否则增强太强CIFAR 这类简单数据集会掉点。下面是完整的预处理代码直接用在DataLoader里import torchvision.transforms as T from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader # 训练增强 train_transform T.Compose([ T.RandomResizedCrop(224, scale(0.3, 1.0)), # 小目标多就调高下界 T.RandomHorizontalFlip(p0.5), T.RandAugment(num_ops2, magnitude9), # num_ops 是每次应用的增强数 T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) # 验证增强 val_transform T.Compose([ T.Resize(256), # 先放大到 256再中心裁剪到 224保留更多上下文 T.CenterCrop(224), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) train_dataset ImageFolder(rootdataset/train, transformtrain_transform) val_dataset ImageFolder(rootdataset/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers8, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers8, pin_memoryTrue) print(类别映射:, train_dataset.class_to_idx)代码逻辑说明RandomResizedCrop(224, scale(0.3, 1.0))表示从原图随机裁剪一块面积为原来的 30% 到 100% 的区域然后缩放到 224x224。scale 下界调高意味着裁剪区域更大目标被截断的概率降低。RandAugment的num_ops决定每次抽样几个增强操作旋转、平移、颜色抖动等magnitude决定增强强度。ImageFolder会根据子目录名称自动生成索引print出来核对 class.json 是否一致。3.3 类别不均衡CrossFormer 的 loss 怎么调自定义数据集经常出现类别不均衡比如class.json里 10 个类其中一个类占了 60% 的样本。CrossFormer 本身不会自动处理这个问题需要你在损失函数层面做调整。最直接的做法是用带权重的CrossEntropyLoss权重和样本数成反比。怎么计算权重有一个常见做法对于每个类别 i权重设为total_samples / (num_classes * class_samples[i])。这样多数类权重接近 1少数类权重大于 1。我在实际项目中更推荐用 Label Smoothing尤其是在类别混淆度高的情况下。CrossFormer 的强表达能力容易让模型对训练集过度自信label smoothing 能有效抑制这个问题。下面是两种 loss 的实现对比from torch.nn import CrossEntropyLoss import torch # 方案一class weight 平衡 def make_class_weight(class_samples, total_samples, num_classes): # class_samples: list每个类别的样本数 # 样本少的类别得到更大的权重 weights [total_samples / (num_classes * n) for n in class_samples] return torch.tensor(weights, dtypetorch.float32) class_weight make_class_weight([1000, 300, 80], 1380, 3) loss_fn_weighted CrossEntropyLoss(weightclass_weight) # 方案二label smoothing 平滑 loss_fn_smooth CrossEntropyLoss(label_smoothing0.1) # 实际训练时二选一或者两者叠加 # 我一般不均衡严重用 weight均衡但困难样本多就用 label smoothing代码逻辑说明CrossEntropyLoss的weight参数要求是一个和类别数等长的 tensor值越大表示该类别的损失权重越高。label_smoothing0.1会把 one-hot 标签变成 0.9 和 0.1/num_classes 的组合防止模型输出极端概率。两个方案可以叠加但注意叠加后学习率要调低否则 loss 前期下降会变慢这是正常现象。4. 训练脚本拆解从 optimizer 到 cosine schedule 的完整配置4.1 AdamW 与余弦退火CrossFormer 训练的三个关键参数CrossFormer 作为类 ViT 架构训练时对优化器和学习率的要求比 CNN 高。我不推荐直接上 SGD虽然它在某些 CNN 上表现好但 ViT 类模型用 AdamW 收敛更稳、对初始学习率不敏感。官方实现里默认用 AdamW与 ViT 系模型的训练经验一致。这里讲一下三个关键参数的作用lr、weight_decay、warmup_steps。lr初始值对 224 输入、batch size 64 的规模来说3e-4 到 5e-4 是比较安全的区间。如果你的 batch size 是 256线性缩放 lr 到 1e-3 左右。weight_decay在 ViT 系模型里通常设 0.05但自定义数据集上我建议从 0.01 开始过大的 weight decay 会让小数据集欠拟合。warmup_steps一般占总训练步数的 5% 到 10%因为 Transformer 在训练初期很不稳定没有 warmup 直接大步长更新loss 容易炸掉。下面是一份可以复制运行的训练脚本骨架覆盖了 epoch 循环、验证和 checkpoint 保存import torch import torch.optim as optim from torch.cuda.amp import GradScaler, autocast from tqdm import tqdm def train_one_epoch(model, loader, optimizer, criterion, scaler, epoch): model.train() total_loss, correct, total 0, 0, 0 pbar tqdm(loader, descfEpoch {epoch}) for images, labels in pbar: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() with autocast(): # 混合精度训练 outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() _, preds outputs.max(1) total labels.size(0) correct preds.eq(labels).sum().item() total_loss loss.item() * labels.size(0) pbar.set_postfix({loss: loss.item(), acc: 100.0 * correct / total}) return total_loss / total, 100.0 * correct / total model model.cuda() optimizer optim.AdamW(model.parameters(), lr4e-4, weight_decay0.01) criterion torch.nn.CrossEntropyLoss(label_smoothing0.1) scaler GradScaler() # 余弦退火 warmup 的常用实现 total_steps len(train_loader) * 100 # 假设训练 100 epoch warmup_steps int(total_steps * 0.05) def lr_lambda(step): # step warmup_steps 时线性上升之后按余弦衰减 if step warmup_steps: return (step 1) / warmup_steps progress (step - warmup_steps) / max(1, total_steps - warmup_steps) return 0.5 * (1 torch.cos(torch.tensor(progress * 3.1415926))) scheduler optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) for epoch in range(1, 101): train_loss, train_acc train_one_epoch(model, train_loader, optimizer, criterion, scaler, epoch) scheduler.step() # 每个 epoch 保存一次 torch.save({ model: model.state_dict(), optimizer: optimizer.state_dict(), epoch: epoch, acc: train_acc, }, fcheckpoint_epoch_{epoch}.pth)代码逻辑说明autocast开启混合精度GradScaler负责动态缩放 loss避免 fp16 下梯度下溢。lr_lambda里 warmup 阶段学习率从 0 线性增加到 1 倍之后按余弦曲线从 1 衰减到接近 0注意 LambdaLR 的step在每个 batch 后调用scheduler.step()时递增。保存 checkpoint 时把 optimizer 和 epoch 都存下来方便断点续训。如果你的显存不够把 batch size 减半、gradient accumulation 为 2 即可效果等价但多花一点时间。4.2 梯度累积与大 batch 模拟小显存跑 CrossFormer 的折中方案CrossFormer-S 的参数量大约在 30M 级别224 输入下 batch size 64 需要 12GB 左右的显存。如果你只有 8GB 显存就需要梯度累积来模拟大 batch。具体做法是每走 N 个 batch 才做一次 optimizer.step()loss 做平均。我把常见做法写成代码这样就不用重复造轮子accumulation_steps 2 # 模拟 batch_size 64 * 2 128 optimizer.zero_grad() for i, (images, labels) in enumerate(train_loader): images, labels images.cuda(), labels.cuda() with autocast(): outputs model(images) loss criterion(outputs, labels) / accumulation_steps # 平均防止梯度爆炸 scaler.scale(loss).backward() if (i 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()代码逻辑说明loss / accumulation_steps是为了让累积的梯度和真实大 batch 的梯度尺度一致。注意 BN 类的批归一化在累积时行为会不一样但 CrossFormer 用的是 LayerNorm不受 batch size 影响这也是 ViT 系模型在梯度累积时的一个优势。推荐 accumulation_steps 不要超过 4否则训练速度太慢而且收敛可能出现波动。4.3 迁移学习用 ImageNet 预训练权重还是从零训练我的建议是除非你的数据集和 ImageNet 差异极大比如医学影像否则永远用预训练权重。CrossFormer 的预训练 checkpoint 在官方仓库可以下载加载方式和 ViT 一样但要特别注意最后一个分类头的类别数不同。加载时忽略head层的权重其他层全部加载。下面这段代码处理了权重迁移import torch def load_pretrained(model, ckpt_path, num_classes): # 加载预训练权重忽略分类头 state_dict torch.load(ckpt_path, map_locationcpu) if model in state_dict: state_dict state_dict[model] # 删除 head 层的权重分类头 keys_to_remove [k for k in state_dict.keys() if k.startswith(head.)] for k in keys_to_remove: state_dict.pop(k) # 加载剩余权重strictFalse 允许缺失 head 层 missing_keys, unexpected_keys model.load_state_dict(state_dict, strictFalse) print(缺失键预期只有 head:, missing_keys) print(多余键:, unexpected_keys) return model # 使用示例 model CrossFormer(num_classes10, depths[2,2,6,2], dims[64,128,256,512]) model load_pretrained(model, crossformer_s.pth, num_classes10)代码逻辑说明strictFalse是关键它允许待加载的 state_dict 和模型不完全对齐。删除head.前缀的键是因为预训练模型分类头输出 1000 类而你的数据集只有 10 类。加载后权重被冻结不会自动发生默认全部可训练这只是初始化手段。如果你想微调时只训练 head 层、冻结 backbone需要手动设置 requires_gradFalse但一般数据集没那么少时没必要冻结。5. 避坑指南CrossFormer 训练中的五个典型翻车现场5.1 现象换分辨率后推理直接报 shape mismatch 错误原因CrossFormer 的 RPE 表和输入分辨率强相关。你训练时用 224x224推理时直接扔一张 384x384 的图进去LSA 和 SDA 的注意力矩阵尺寸变了但 RPE 表没有跟着变于是报错。解决回到 2.3 节的reset_rpe_for_resolution函数在切换分辨率后重建 RPE 表。注意重建后需要重新跑几个 epoch 微调否则精度会下降。血泪经验我一开始写好模型后直接做一次推理测试发现报错以为 bug实际是 RPE 没按分辨率重置。5.2 现象训练 loss 收敛到 1.0 左右就下不去原因RandomResizedCrop的 scale 下界设太低配合RandAugment的强增强CrossFormer 学到的特征被增强噪声干扰。另一个常见原因是 label smoothing 设得过大比如 0.3会限制 logit 的置信度上限loss 自然降不到很低。解决把 scale 下界从 0.08 调到 0.3RandAugment的 magnitude 从 9 降到 5label smoothing 保持 0.1。如果还是降不下去检查 class.json 里类别数是否和模型num_classes一致有时类别数写多了模型一直在预测不存在的类。5.3 现象训练时 GPU 显存溢出原因batch size 太大或者输入分辨率太高。CrossFormer 的 SDA 在下采样倍数较小时显存占用会显著上升。尤其当subsample_factor设为 1等价于全分辨率全局注意力显存占用直接爆炸。解决用 4.2 节的梯度累积替代硬扛 batch size。另外检查subsample_factor列表深层 stage 不要设 1除非显存充足。我之前把 stage 4 的subsample_factor设成 1 想提升精度结果 24GB 显存的卡直接 OOM改回 2 就稳了。5.4 现象迁移学习后模型精度比从零训练还低原因预训练模型是在 ImageNet 上过拟合过 1000 类分布的如果你的数据分布差异大微调时学习率过高会把预训练特征冲掉导致「灾难性遗忘」。更隐蔽的原因是预训练权重的 RPE 表是和 224 分辨率绑定的你如果没重置直接微调注意力位置信息全是错的。解决微调时把 lr 降到 1e-4 以下前 5 个 epoch 只训练 head 层冻结 backbone之后再解冻全部参数用低学习率训练。另外检查是否执行了 RPE 重置。从那以后我每次迁移前都会强制走一遍这两个流程基本没再翻过车。5.5 现象推理时输出概率全是接近均匀分布原因模型最后是用nn.Linear做分类头如果num_classes和实际类别数不匹配但训练时没有报错因为数据加载器只用了实际存在的类别那么 head 层有大量神经元从来没有梯度更新推理时它们输出随机噪声拉低置信度。解决用 3.1 节的方式显式打印train_dataset.class_to_idx比对 class.json 的映射。如果发现不一致用映射后的索引重建 class.json再重新训练。这种问题往往在自定义数据集上出现因为目录结构看似正确实际类别映射已经被 ImageFolder 的字典序打乱了。6. 进阶验证技巧把 LSA 和 SDA 的注意力图抽出来看看模型训完精度达标这只是第一步。要确认 CrossFormer 是真的学会了跨尺度交互还是单纯在过拟合我会把 LSA 和 SDA 的 attention map 抽出来做可视化。这个操作比只盯着 accuracy 有用得多——它能直接告诉你模型在关注哪些区域以及小尺度和大尺度的信息是怎么融合的。具体做法在前向传播时注册 hook把 LSA 和 SDA 模块的 attention 权重保存下来。SDA 的 attention 权重形状是[B, num_heads, HW/L^2, HW/L^2]取第一行就是第一个 token 对所有 token 的注意力分布。将其 reshape 成特征图分辨率插值到原图大小再叠加到原图上就能看到每个像素对最终分类的贡献权重。工具代码不复杂核心是要拿到 attention 权重而不是输出特征。用下面这段代码可以快速实现import torch import matplotlib.pyplot as plt import torch.nn.functional as F attention_maps {} def hook_fn(name): def forward_hook(module, input, output): # 从模块的输出中抽取 attention 权重 if hasattr(module, attn): attention_maps[name] module.attn.detach().cpu() return forward_hook # 注册 hook找到所有 LSA 和 SDA for name, module in model.named_modules(): if lsa in name or sda in name: module.register_forward_hook(hook_fn(name)) def visualize_attention(model, img_tensor, save_pathattention.png): model.eval() with torch.no_grad(): _ model(img_tensor.unsqueeze(0)) # 取最后一个 stage 的 SDA 注意力 for name, attn in attention_maps.items(): if sda in name and 3 in name: # 第 4 个 stage 的 SDA # attn: [B, num_heads, HW, HW] attn_map attn[0, 0] # 取第一个 head h w int(attn_map.shape[0] ** 0.5) # 假设是正方形 attn_map attn_map.reshape(h, w) attn_map F.interpolate( attn_map.unsqueeze(0).unsqueeze(0), size(224, 224), modebilinear ).squeeze() # 归一化到 0-1 attn_map (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min()) # 叠加原图 img_np img_tensor.permute(1, 2, 0).numpy() img_np (img_np * 0.5 0.5) # 反归一化 plt.figure(figsize(8, 4)) plt.subplot(1, 2, 1) plt.imshow(img_np) plt.title(原图) plt.subplot(1, 2, 2) plt.imshow(img_np) plt.imshow(attn_map.numpy(), cmapjet, alpha0.5) plt.title(SDA 注意力热力图) plt.savefig(save_path, dpi150, bbox_inchestight) print(f可视化结果保存到 {save_path})代码逻辑说明hook 函数在 forward 执行完后触发把module.attn取出来这个属性是模块内部保存的注意力权重矩阵。F.interpolate将低分辨率的 attention map 上采样到原图大小cmapjet叠加半透明彩色热力图红色代表高关注区域。实际操作中你会发现训练好的 CrossFormer 的 SDA 热力图会同时覆盖目标主体和部分上下文背景这正是跨尺度注意力希望在全局和局部之间取得平衡的表现。如果热力图只集中在一个小点或者完全发散没有重点说明模型欠拟合或数据增强过强需要回来调参。这个习惯帮我在两个项目里发现了数据标注错误热力图高亮区不在目标物体上检查后发现是图片和标签对错了。从那以后我每次换数据集训练完都会强制走一遍这个可视化流程再写报告。希望帮到你。本文还有配套的精品资源点击获取