★ FEATURED ARTICLE Day 47我把一整天都耗在了“注意力热图可视化”上。原本以为就是跑几行代码、存几张图的事结果从模型中间层抽权重、对齐 token 和像素坐标、归一化、多头上采样一路踩坑踩到晚上。等真的看到热图那一刻还是值了——模型在哪看、没在看什么一目了然。注意力热图可视化简单说就是把 Transformer 这类模型内部的注意力权重从张量里取出来用颜色映射到原图上告诉人“模型处理这个位置时把计算资源分配给了哪些区域”。适合做模型解释、论文配图、Debug 坏 case、判断模型是不是学偏了也适合刚接触可解释性的同学用来建立直觉。这篇文章不聊太多理论主要是我自己的实操路线、代码片段和一批很值钱的坑。1. 注意力热图到底在画什么先搞懂可视化背后的逻辑1.1 注意力权重是什么模型内部生成的一张“关系表”要画热图首先得明确一件事注意力热图画的不是模型“敢不敢看”某个地方而是模型在处理某个 token 时对其他 token 分配了多少权重。这个权重来自注意力机制里的自注意力矩阵。用大白话说注意力机制就是模型在“开大会”。每个 token 或者图像 patch 先拿一张“查询牌”Q向全场提问其他 token 拿着自己的“键信息”K 回应系统做一个打分匹配后再经过 softmax 变成一组和为 1 的权重分数最后按这个权重从“值信息”V 里带回内容。所以模型在每一层、每一个注意力头上都会产出一个形状为 [序列长度, 序列长度] 的矩阵。这个矩阵就是一张“关系表”行表示“我是谁在查询”列表示“我看了谁”。热图可视化的本质就是把这张关系表中的某一行、某一列或者一个聚合结果映射成不同颜色画出来。红色的区域对应高权重蓝色区域对应低权重直观得不行。但很多人第一步就卡住了模型那么多层、那么多头到底取哪一层、哪个头、哪一行这不是随意选的后文会详细说。1.2 注意力热图和 CAM 不是一回事两套思路别混用我见过不少同学把注意力热图和 Grad-CAM 当成同一种东西其实两者的原理和回答的问题完全不同。Grad-CAM 这类基于梯度的可视化方法需要你给出类别标签然后计算“类别置信度对特征图的梯度”再对梯度做全局平均得到每个通道的重要性权重最后和特征图加权求和。它回答的问题是模型为了将这张图认成某个类别用到了哪些区域的证据注意力热图回答的问题则是模型在处理 CLS token 时把注意力分配给了哪些 patch它不需要标签不需要梯度只要在模型前向传播时顺手把注意力权重存下来就行。一个是梯度的产物一个是权重的产物语义解释差别很大。实操里这两者常常放在一起做对比展示。比如我做过一个端侧分类项目某张坏样本的 Grad-CAM 显示模型用了狗头而注意力热图还同时瞄着背景里的草坪。这时就能得到更多信息模型不是不知道狗头在哪而是背景信息也参与了决策。这种组合分析比单独画一种图更能定位问题。1.3 热图给谁看论文配图、模型诊断与业务解释我也总结了注意力热图最常见的三种用途方便你判断自己到底需要哪种画法。第一展示与汇报。做学术项目或者给业务方交付的时候一张“原图和热图叠加”的图比失准的准确率数字更有说服力。这个场景要求图好看、颜色语义清晰、分辨率高输出尺寸一般要按发表或者打印标准来。第二模型诊断。这类用途最吃细节。模型在某个验证子集上一直犯错把热图叠加到测试图上一眼就能看出它是不是把视觉重心放在了水印、拍摄时间戳、甚至是背景物体上。如果模型通过“捷径”来分类注意力会非常集中在一个不该关注的位置这个时候热图就是最好的证据。第三数据筛查。批量跑一批样本的热图去找“模型注意力与人工标注区域不匹配”的样本。这些样本要么是标注错了要么是图像本身有歧义自动筛出来之后人工二次确认能省下大量无脑刷数据的时间。这一用途对代码的批量处理能力有要求后面会讲到怎么做一个极简的批处理脚本。2. 拿到注意力权重的三条路径模型不同方法完全不一样2.1 transformers 库模型最省心output_attentions 直接拿如果你用的是 transformers 体系下的 ViT、BERT 这类模型拿注意力权重几乎不费脑子。只需要在调用接口的时候显式打开output_attentions输出的对象里就会多出一个attentions字段里面保存每一层的注意力矩阵。import torch from transformers import AutoModel, AutoImageProcessor from PIL import Image # 用你本地或远程仓库里的权重路径比如 vit-base-patch16-224 model_path vit-base-patch16-224 model AutoModel.from_pretrained(model_path, output_attentionsTrue) model.eval() processor AutoImageProcessor.from_pretrained(model_path) image Image.open(demo.jpg).convert(RGB) inputs processor(imagesimage, return_tensorspt) with torch.no_grad(): outputs model(**inputs, output_attentionsTrue) # attentions 是一个元组每一层对应一个 [batch, heads, seq_len, seq_len] 的张量 attn_layer_6 outputs.attentions[5] print(attn_layer_6.shape) # [1, 12, 197, 197]这里有几个经验点。model.eval()必须加不然模型内部的 dropout 和 BatchNorm 行为不一致拿到的注意力带噪声。torch.no_grad()一定包住前向否则计算图会占额外显存也会拖慢速度。output_attentionsTrue这个参数也可以在from_pretrained之后通过model.config.output_attentions True设置但每次调用时显式传参更不容易踩到缓存 config 的坑。2.2 自定义模型侧录forward_hook 与改源码如果没有用封装好的接口而是自己搭的 ViT、或者复用了某些 GitHub 上的实现拿权重就会稍微麻烦一些。很多简化实现为了省显存forward 的返回值里只有 logits没有 attention weights。我常用的方式是在目标注意力模块上注册一个 forward 钩子。from torch import nn attn_cache {} def make_hook(name): def hook_fn(module, input_, output): # 前提该模块内部把 softmax 后的注意力权重保存到了某个属性里 # 例如 module.attn_weights weights getattr(module, attn_weights, None) if weights is not None: attn_cache[name] weights.detach().cpu() return hook_fn # 遍历模型模块找到注意力模块并注册钩子 for name, module in model.named_modules(): if attention in name.lower() and hasattr(module, attn_weights): module.register_forward_hook(make_hook(name)) # 随后正常做一次前向核心思路是很多实现的attn_weights是作为模块内部中间变量存在的forward 结束后还没被销毁正好在钩子里把它截留。detach().cpu()可以防止梯度累积和显存堆积。如果这个实现没有保存中间变量那就只能改源码了在注意力 forward 的softmax之后加一行self.attn_weights attn_weights.detach()基本是零副作用。这是最土最笨的方法但也是兼容性最高的。2.3 新手最容易踩的点softmax 前后的权重不是一回事画热图前先确认你拿的是 softmax 之前的还是之后的权重。绝大多数情况我们默认使用 post-softmax 权重也就是每一行加起来为 1 的那个版本因为它本身就是归一化的直接画热图不会出现极端大值。但分析某些特定问题时pre-softmax 权重反而更有用。比如你想看某个注意力头是不是“饱和”了softmax 之后所有值挤到 0 和 1 两端已经区分不出细微差异用 pre-softmax 就能看到原始打分。再比如你要做梯度分析softmax 之后的梯度已经被链式法则揉过一遍不如看原始打分那边干净。从实现上看一档的 transformers 库输出的是 post-softmax 权重不用你再算自定义实现里则百分之百要检查一下在哪个位置存的。如果画出来的热图所有 patch 都是几乎相同的浅色、看不出重点优先怀疑你拿的是 softmax 之前的权重而不是图像或者插值的问题。3. 从权重到热图完整实操流程与参数选择3.1 特征索引还原token 到像素的换算关系权重拿出来了下一步就是把它变成空间网格。以 ViT-B/16 为例输入是 224×224 的图像patch size 是 16×16所以 token 网格是 14×14加上一个分类用的 CLS token序列长度就是 197。注意力矩阵最后一个维度的索引 0 到 196 中索引 0 是 CLS索引 1 到 196 按行优先顺序对应 14×14 网格上的 patch。目标如果只是看“CLS token 最后关注了哪些 patch”那就取注意力矩阵中 CLS 那一行去掉第 0 列CLS 对自己的注意力一般很高但画图时通常不看剩下 196 个数直接 reshape 成 14×14 的网格图。# attn: [1, 12, 197, 197] attn outputs.attentions[5][0] # 取出 batch 维度变成 [12, 197, 197] avg_heads attn.mean(dim0) # 多头平均变成 [197, 197] cls_row avg_heads[0, 1:] # CLS 对其他 token 的注意力 grid cls_row.reshape(14, 14) # 行优先展开正好对应图像网格如果用的是文本模型比如 BERT那 token 就是一维的不存在 reshape 的问题。常见做法是画一个 [seq_len, seq_len] 的注意力矩阵热图横纵轴是 token 序号颜色深浅表示 token 之间的注意力大小。这个用 matplotlib 或者 DataFrame 的样式渲染都能快速实现。3.2 渲染设置颜色映射、缩放方式、模糊半径怎么选拿到 14×14 的小矩阵后直接imshow出来会发现锯齿严重、patch 边界生硬。正确做法是先放大到原图大小再做高斯模糊最后叠加显示。我的标准渲染流程如下import cv2 import numpy as np import matplotlib.pyplot as plt from PIL import Image def render_attention(image_path, attn_map_2d): image Image.open(image_path).convert(RGB) w, h image.size # 用双立方插值把 14x14 放大到原图尺寸 attn_up cv2.resize( attn_map_2d.astype(np.float32), (w, h), interpolationcv2.INTER_CUBIC, ) attn_up cv2.GaussianBlur(attn_up, (9, 9), 0) # 用百分位截断而不是全局最大值避免个别高亮点把颜色拉爆 v_max np.percentile(attn_up, 99) v_min attn_up.min() attn_norm np.clip((attn_up - v_min) / (v_max - v_min 1e-8), 0, 1) fig, axes plt.subplots(1, 2, figsize(12, 6)) axes[0].imshow(image) axes[0].axis(off) axes[1].imshow(image) im axes[1].imshow(attn_norm, cmapjet, alpha0.5) axes[1].axis(off) fig.colorbar(im, axaxes[1], fraction0.046, pad0.04) plt.tight_layout() plt.savefig(result.png, dpi150)这里有两个细节非常影响最终效果。第一颜色映射方面jet在展示类的工作里出图最有冲击力红到蓝对比鲜明viridis更友好适合打印和色盲读者。第二透明度 alpha 建议在 0.4~0.6 之间太低了重点不明显太高则盖住原图结构反而看不出注意力区域和物体的对应关系。3.3 多头注意力合并策略平均、最大还是逐头观察ViT 有 12 个注意力头BERT-large 甚至有 16 个这决定了你不可能只用一种方式把多头信息压成一张图。最偷懒的做法是把所有头取平均实际操作中我也经常这么干它适合快速看个大概。但弊端非常明显不同的头可能关注完全不同的语义有的盯轮廓有的盯颜色有的盯背景和物体的位置关系平均之后这些模式互相抵消热图变成一团没有重点的雾。更稳妥的流程是先把所有头单独画在一个网格图里逐头扫一眼找到和任务直觉最吻合的那个头再以它为准出单张大图。比如分类任务里某个头可能把 CLS 的注意力几乎全放在了目标物体上另一些头则在背景上到处晃。你选前者做展示说服力比盲平均强得多。还有一种折中方案适合给非技术背景的人看先把多头平均再保留注意力值排名前 20% 的区域其他全部置零最后再上采样。这样做出来的热图区域性强、抗噪好代价是会丢失一些小权重区域的细节。三个方案我都在不同交付场景里用过“逐头看 选代表”最严谨“平均 阈值”最省事。4. 高频踩坑与排查实录这些坑值得记两晚4.1 尺寸对不上B×H×N×N 如何正确 reshape 回网格第一个高频坑就是把[batch, heads, seq_len, seq_len]张量处理成单层可视化时维度顺序搞乱。我见过不少人直接拿attentions[5]的 shape 去 reshape出来一张斜的、镜像的、或者整体平移的热图。原因很简单你少取了 batch 维或者把 query 维和 key 维搞反了。记住attn[b, h, i, j]中i是查询 token 的索引j是被关注的 key token 索引。看“CLS 关注哪里”要取的是[0, 0, :, :]中的i0行即第 0 个 token 对其他 token 的权重。如果误取了[:, :, :, 0]列你就会画出“其他 token 如何关注 CLS”的图语义完全不同。reshape 成 14×14 时还要确认你的 token 顺序是行优先扫过网格的。绝大多数实现都满足这个假设但个别修改过的实现可能先列后行或者用了不同的 patch 展开顺序。验证办法很简单取网格左上角几个 patch 的索引值和图像左上角区域的特征做对应检查肉眼比对热图和原图就能发现问题。4.2 热图总是糊的放大插值与高斯模糊的正确姿势热图发糊是第二个高频问题而且原因往往不是渲染代码而是注意力权重本身。ViT 浅层的注意力通常比较均匀每个 patch 都会被分配到差不多的权重用线性插值放大再一糊整张图就是一片渐变。这时候直接把权重贴上去没意义建议先做归一化或者只选深层比如第 6 层之后来可视化。反过来如果你刻意想保留 patch 之间的边界比如要展示“模型精确聚焦到了某 2×2 的 patch 区域”就别用双立方插值改用INTER_NEAREST然后用大一点的模糊核稍微润色。我实测下来高斯模糊核大小与放大倍率有关14×14 放大到 224×224 时核选 7×7 到 11×11 比较舒服sigma 可以直接给 0让它自己按核尺寸算。4.3 Swin 这类局部注意力模型怎么画别照搬全局 token 的逻辑Swin Transformer 这类使用窗口注意力的模型画热图和 ViT 完全不是一回事。它的自注意力不是在整个 14×14 网格上做的而是在一个个局部窗口内部做的。也就是说一个 token 的注意力范围只覆盖它所在窗口内的几十个 token拿到的注意力矩阵不能直接 reshape 成全局网格。如果你只是“看一眼效果”又不想在 Swin 的内部实现里挣扎太久我的建议是直接换成 CAM 或 Grad-CAM 可视化。这不算偷懒而是工具选型问题局部注意力模型天然就不适合做全局注意力热图硬画出来的窗口碎片不仅误导人还容易让业务方觉得“模型注意力很局限”。如果确实要画方向是按窗口索引还原# 伪代码从 Swin 窗口注意力还原为全局特征图 # 1. 拿到窗口内注意力: shape [batch * window_num, heads, window_size^2, window_size^2] # 2. 按窗口位置 scatter 回 feature map 的空位 # 3. 处理 shifted window 时需要还原之前 roll 过的偏移最后一行的 shifted window 处理是最容易出错的地方。Swin 通过循环平移改变窗口划分如果不做反向的 roll 操作还原出来的网格整体会错位。真要做这一步建议对照源码逐步调试别直接套网上的现成代码。4.4 加了 output_attentions 就 OOM显存和性能问题很多人会遇到一个情况不开output_attentions模型能跑一开就显存溢出。因为注意力矩阵的大小是序列长度的平方还乘上头数和层数。对于 ViT 这种 197×197 的规模还好但如果是一个长文本模型序列长度到 2048、4096那个矩阵就很可观了。我的处理逻辑是不要一次性保留所有层的注意力。先用一个很短的 demo 输入确认通道没问题然后只对目标层做一次前向拿完马上del和清缓存。如果用钩子截流限定只保存第一个 batch 样本的权重不要整张attn_weights都往列表里塞不然显存和内存一起炸。另外调试阶段建议用小尺寸输入跑通确认逻辑没问题再上真实分辨率可以省下不少时间。5. 热图之外的三个延伸方向把可视化变成调试工具5.1 注意力流逐层追踪信息传播路径单层的注意力热图只告诉你某一层怎么分配注意力但信息是层层传播的。把相邻层的注意力权重连乘起来可以近似模拟信息从浅层到深层的流动路径也就是论文里常说的 Attention Rollout。方法很简单从第一层开始逐层做权重矩阵乘法得到深层 token 实际“溯源”到浅层 token 的累计权重。实操时要注意加上恒等映射也就是把每层的注意力矩阵和单位矩阵加权融合再连乘。否则多层连乘后注意力衰减得极快深层几乎看不出差异。这个技术在做模型压缩时特别有用如果某些注意力流的路径非常固定就说明有冗余头可以剪掉。5.2 从定性到定量注意力熵与聚焦度指标热图看多了之后你会发现人眼判断太主观。同样是“比较集中”有的模型集中往返于两个区域有的则死死盯在一个 patch 上。为了把这种感觉变成可比较的指标我会额外算两个数注意力熵和 top-1 注意力占比。注意力熵就是把每个 token 行的注意力分布当成概率分布算信息熵。熵越低说明这一行的注意力越集中在少量 token 上模型越“笃定”。top-1 注意力占比是每行最大的那个注意力值它能反映聚焦程度。这两个指标可以在模型训练的不同 checkpoint 上追踪用来判断训练是否稳定比肉眼看热图可量化得多。一张热图配合一组指标解释报告的可信度立刻高一档。5.3 输出整合批量生成解释报告最后一个延伸方向就是把单张热图变成批量报告。做法不复杂准备一个文件夹的测试图循环走一遍“前向 - 取权重 - 渲染”把每张图的“原图 热图 预测类别 置信度”拼到一张大图里最后存成 HTML 或者 PDF。我习惯存 HTML因为可以插入 hover 事件鼠标移到图上直接看对应 patch 的注意力数值。这个批量脚本最大的价值是在模型交付前做一次全面体检哪些样本模型注意力特别分散、哪些样本置信度高但注意力在背景上都能一次性拉出来。配合误分类样本很快就能找到数据集的脏点。Day 47 做下来我最大的体会不是“我会画热图了”而是“可视化是最容易骗人的工具”。热图能帮你发现问题也能让你误以为模型很聪明。在你把一张热图当结论之前至少应该看一眼不同层、不同头、不同样本确认它不是偶然。养成了这个习惯注意力热图才真正变成你的一个调试工具而不是一张放在报告里好看的花图。
阅读完成 · 觉得有帮助？
