最强辅助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_localdef 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