最强辅助Visualizer：简化你的Vision Transformer可视化！

©作者 | 洛英

Visualizer 是一个辅助深度学习模型中 Attention 模块可视化的小工具，主要功能是帮助取出嵌套在模型深处的 Attention Map。

Vision Transformer 如今已经成为了一个热门的方向，Self-Attention 机制为视觉信息的表示和融合都带来了新的思路。那么问题来了，作为 Attention 机制的一种，Self-Attention 会呈现出怎样的一种 Q K V 交互模式呢？为了搞清这个问题，我们只要将 Vit 中间的 Attention Map 拿出来看看不就行了嘛，然而把 Attention Map 拿出来可不是那么简单的事情。

为了可视化 Attention Map，你是否有以下苦恼：

1. Return 大法好：通过 return 将嵌套在模型深处的 Attention Map 一层层地返回回来，然后训练模型的时候又不得不还原；

2. 全局大法好：使用全局变量在 Attention 函数中直接记录 Attention Map，结果训练的时候忘改回来导致 OOM。

不管你有没有，反正我有，由于可视化分析不是一锤子买卖，实际过程中你往往需要在训练-可视化-训练-可视化两种状态下反复横跳，所以不适合采用以上两种方式进行可视化分析。

PyTorch hook 的局限性

咨询了专业人士的意见后，发现 pytorch 有个 hook 可以取出中间结果，大概查了一下，发现确实可以取出中间变量，但需要进行如下类似的 hook 注册。

handle = net.conv2.register_forward_hook(hook)

这样我们就可以拿出来 net.conv2 这层的输出啦。

然而！进行这样操作的前提是我们知道要取出来的模块名，但是 Transformer 类模型一般是这样定义的（以 Vit 为例）。

class VisionTransformer(nn.Module):
    def __init__(self, *args, **kwargs):
        ...
        self.blocks = nn.Sequential(*[Block(...) for i in range(depth)])
        ...

然后每个 Block 中都有一个 Attention 。

class Block(nn.Module):
    def __init__(self, *args, **kwargs):
        ...
        self.attn = Attention(...)
        ...

然后我们想要的 attention map 又在 Attention 里面。

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        ...

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn) # <-在这

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

如果想要用 hook 来取出 Vit 中的 Attention Map，存在的问题就是：

1. 嵌套太深，模块名不清晰，我们根本不知道我们要取的 attention map 怎么以 model.bla.bla.bla 这样一直点出来！ 

2. 一般来说，Transformer 中 attention map 每层都有一个，一个个注册实在太麻烦了。

那怎么办呢....

Visualizer！ 

所以我就思考并查找能否通过更简洁的方法来得到 Attention Map（尤其是 Transformer 的），而 visualizer 就是其中的一种，它具有以下特点：

• 精准直接，你可以取出任何变量名的模型中间结果；

• 快捷方便，一个操作，就可以同时取出 Transformer 类模型中的所有 attention map；

• 非侵入式，你无须修改函数内的任何一行代码；

• 训练-测试一致，可视化完成后，训练时无须再将代码改回来。

用法

项目主页：

https://github.com/luo3300612/Visualizer

首先，git clone 并安装它：

python setup.py install

使用方法非常简单！只需要用 get_local 装饰一下 Attention 的函数，forward 之后就可以拿到函数内与装饰器参数同名的局部变量。

使用方法一

比如说，我想要函数里的 attention_map 变量。在模型文件里，我们这么写：

from visualizer import get_local

@get_local('attention_map') # 我要拿attention_map这个变量,所以把他传参给get_local
def your_attention_function(*args, **kwargs):
    ...
    attention_map = ... 
    ...
    return ...

在可视化代码里，我们这么写：

from visualizer import get_local
get_local.activate() # 激活装饰器
from ... import model # 被装饰的模型一定要在装饰器激活之后导入！！

# load model and data
...
out = model(data)

cache = get_local.cache # ->  {'your_attention_function': [attention_map]}

最终就会以字典形式存在 get_local.cache 里，其中 key 是你的函数名，value 就是一个存储 attention_map 的列表。

使用方法二

使用 Pytorch 时我们往往会将模块定义成一个类，此时也是一样只要装饰类内计算出 attention_map 的函数即可：

from visualizer import get_local

class Attention(nn.Module):
    def __init__(self):
        ...

    @get_local('attn_map')
    def forward(self, x):
        ...
        attn_map = ...
        ...
        return ...

其他细节请参考：

https://nbviewer.jupyter.org/github/luo3300612/Visualizer/blob/main/demo.ipynb

可视化结果

这里是部分可视化 vit_small 的结果，全部内容在 demo.ipynb 文件里。

因为普通 Vit 所有 Attention map 都是在 Attention.forward 中计算出来的，所以只要简单地装饰一下这个函数，我们就可以同时取出 vit 中 12 层 Transformer 的所有 Attention Map！ 

一个 Head 的结果：

一层所有 Heads 的结果：

红色 grid 作为 query 的 Attention Map：

有趣的结果

在可视化这张图片的过程中，我也发现了一些有趣的现象。 

首先，靠前层的 Attention 大多只关注自身，进行真·self attention 来理解自身的信息，比如这是第一层所有 Head 的 Attention Map，其特点就是呈现出明显的对角线模式。

随后，模型开始逐渐增大感受野，融合周围的信息，呈现出多条对角线的模式，如下分别是第 4、6 层的 Attention Map。

最后，重要信息聚合到某些特定的 token 上，Attention 出现与 query 无关的情况，在 Attention Map 上呈现出竖线的模式，如下第 11 层的 Attention Map：

当然，这些只是一张图片可视化的结果，说不定多看几张图片还有更多或者不同的结论呢。

注意

在使用 visualizer 的过程中，有以下几点需要注意：

1. 想要可视化的变量在函数内部不能被后续的同名变量覆盖了，因为 get_local 取的是对应名称变量在函数中的最终值；

2. 进行可视化时，get_local.activate() 一定要在导入模型前完成，因为 python 装饰器是在导入时执行的；

3. 训练时你不需要删除装饰的代码，因为在 get_local.activate() 没有执行的情况下，attention 函数不会被装饰，故没有任何性能损失（同上一点，因为 python 装饰器是在导入时执行的）。

其他

当然，其实 get_local 本质就是获取一个函数中某个局部变量的最终值，所以它应该还有其他更有趣的用途。

小结

Visualizer 是一个辅助深度学习模型中 Attention 模块可视化的小工具，主要功能是帮助取出嵌套在模型深处的 Attention Map，欢迎大家多多 star!

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·