Visualizer

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

为什么需要 Visualizer？

为了可视化 Attention Map，你是否有以下困扰：

• “Return 大法好”：通过 return 将嵌套在模型深处的 Attention Map 一层层地返回回来，然后训练模型的时候又不得不还原。
• “全局大法好”：使用全局变量在 Attention 函数中直接记录 Attention Map，结果训练的时候忘改回来导致 OOM。

不管你有没有，反正我有。

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

handle = net.conv2.register_forward_hook(hook)

进行这样操作的前提是我们知道要取出来的模块名。然而，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(...)
        ...

如果要使用 hooks，其中的问题就是：

1. 嵌套太深，模块名不清晰，我们根本不知道我们要取的 attention map 怎么以 model.bla.bla.bla 这样一直点出来！
2. 一般来说，Transformer 中 attention map 每层都有一个，一个个注册实在太麻烦了。

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

• 精准直接，你可以取出任何变量名的模型中间结果。
• 快捷方便，同时取出 Transformer 类模型中的所有 attention map。
• 非侵入式，你无须修改函数内的任何一行代码。
• 训练-测试一致，可视化完成后，你无须在训练时再将代码改回来。

使用方法

安装

pip install bytecode
python setup.py install

安装完成后，只需要用 get_local 装饰一下 Attention 的函数，forward 之后就可以拿到函数内与装饰器参数同名的局部变量啦~

Usage1

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

from visualizer import get_local
@get_local('attention_map')
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 的列表。

Usage2

使用 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 ...

其他细节请参考 demo.ipynb 文件。

可视化结果

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

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

一个 Head 的结果：

一层所有 Heads 的结果：

某个 grid 的 Attention Map：

注意事项

• 想要可视化的变量在函数内部不能被后续的同名变量覆盖了，因为 get_local 取的是对应名称变量在函数中的最终值。
• 进行可视化时，get_local.activate() 一定要在导入模型前完成，因为 Python 装饰器是在导入时执行的。
• 训练时你不需要修改/删除任何代码，即不用删掉装饰函数的代码，因为在 get_local.activate() 没有执行的情况下，attention 函数不会被装饰，故没有任何性能损失（同上一点，因为 Python 装饰器是在导入时执行的）。

其他

当然，其实 get_local 本身可以取出任何一个函数中某个局部变量的最终值，所以它应该还有其他更有趣的用途。

参考文献

• bytecode
• local track1
• local track2
• decorator1
• decorator2

Visualizer 快速上手指南

Visualizer 是一个专为深度学习设计的轻量级工具，旨在非侵入式地提取模型内部（特别是 Transformer 架构中）嵌套深层的 Attention Map。无需修改模型训练逻辑，即可在推理或可视化阶段轻松获取中间变量。

环境准备

在使用前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：>= 3.6
• 核心依赖：
  • PyTorch (版本 > 0.0.0)
  • bytecode

安装步骤

建议先安装底层依赖 bytecode，然后安装本工具。

pip install bytecode
python setup.py install

提示：如果下载速度较慢，可使用国内镜像源加速安装：

pip install bytecode -i https://pypi.tuna.tsinghua.edu.cn/simple
python setup.py install

基本使用

Visualizer 的核心是通过 @get_local 装饰器标记需要提取的变量，并在导入模型前激活该功能。

1. 修改模型代码（仅需添加装饰器）

在你的模型文件中，找到计算 Attention Map 的函数（可以是独立函数或类中的 forward 方法），使用 @get_local('变量名') 进行装饰。

示例：装饰类中的 forward 方法

from visualizer import get_local
import torch.nn as nn

class Attention(nn.Module):
    def __init__(self):
        super().__init__()
        # ... 初始化代码 ...
    
    @get_local('attn_map')  # 标记需要提取的局部变量名为 'attn_map'
    def forward(self, x):
        # ... 计算逻辑 ...
        attn_map = ...  # 确保函数内存在此变量且未被后续同名变量覆盖
        # ... 其余逻辑 ...
        return output

2. 执行可视化脚本

在运行推理或可视化的脚本中，务必在导入模型之前调用 get_local.activate()。

from visualizer import get_local

# 第一步：激活装饰器（必须在导入模型前执行）
get_local.activate()

# 第二步：导入已修饰的模型
from your_model_file import VisionTransformer 

# 第三步：加载模型与数据并运行
model = VisionTransformer()
data = ... # 准备输入数据
out = model(data)

# 第四步：获取结果
# cache 是一个字典：{ '函数名': [变量值列表] }
cache = get_local.cache 

# 提取特定函数的 attention map
attention_maps = cache['forward'] 
print(f"提取到的 Attention Map 数量：{len(attention_maps)}")

关键注意事项

• 导入顺序：get_local.activate() 必须在 import model 之前执行，因为 Python 装饰器在模块导入时即生效。
• 变量命名：确保装饰器参数中的名称（如 'attn_map'）与函数内部最终想要获取的变量名完全一致，且该变量在函数结束前未被重新赋值覆盖。
• 训练无影响：若不执行 activate()，装饰器不生效，原模型代码无任何性能损耗，可直接用于正常训练。