视觉Transformer的可解释性（基于PyTorch）

本仓库实现了视觉Transformer模型的可解释性方法。

更多信息请参见：https://jacobgil.github.io/deeplearning/vision-transformer-explainability

当前已实现的功能：

• 注意力回传（Attention Rollout）。

• 针对特定类别的梯度注意力回传。 这是我们在此基础上进一步改进和优化注意力回传的尝试。

• 待定：注意力流仍在开发中。

此外，还包含一些调整和技巧以确保其正常运行：

• 不同的注意力头融合方法。
• 移除最低权重的注意力。

使用方法

• 通过代码调用

from vit_grad_rollout import VITAttentionGradRollout

model = torch.hub.load('facebookresearch/deit:main', 
'deit_tiny_patch16_224', pretrained=True)
grad_rollout = VITAttentionGradRollout(model, discard_ratio=0.9, head_fusion='max')
mask = grad_rollout(input_tensor, category_index=243)

• 通过命令行调用：

python vit_explain.py --image_path <图片路径> --head_fusion <mean, min 或 max> --discard_ratio <0到1之间的数字> --category_index <类别索引>

如果未指定category_index，则使用注意力回传；否则将使用梯度注意力回传。

请注意，默认情况下，该工具使用的是来自高效数据训练的图像Transformer及注意力蒸馏论文中的“Tiny”模型，该模型托管在PyTorch Hub上。

Transformer在这张图片中关注了哪些区域？

图片	原始注意力回传	带丢弃比例+最大值融合

针对特定类别的梯度注意力回传

Transformer内部流动的注意力会传递不同类别的信息。通过梯度回传，我们可以看到网络关注的具体位置，但它无法告诉我们这些位置是否最终被用于最终的分类决策。

我们可以通过将注意力与目标类别的输出梯度相乘，并在屏蔽掉负值注意力后对所有注意力头取平均值，从而只保留那些对目标类别（或多个类别）有贡献的注意力。

Transformer在哪些地方看到了狗（类别243）和猫（类别282）？

Transformer在哪些地方看到了猎犬（类别161）和鹦鹉（类别87）：

使其实现的一些技巧和调整

过滤每一层中最低权重的注意力

--discard_ratio <0到1之间的值>

通过保留最强的注意力来去除噪声。

不同丢弃比例的效果：

不同的注意力头融合方式

注意力回传方法建议对所有注意力头的注意力进行平均，但根据经验来看，取最小值或结合丢弃比例取最大值的方式效果更好。

--head_fusion <mean, min 或 max>

图片	平均融合	最小值融合

参考文献

• 量化Transformer中的注意力流

• timm：一个优秀的PyTorch模型集合，尤其是其中的视觉Transformer实现

• 一张图胜过16×16个词：大规模图像识别中的Transformer

• 感谢https://github.com/jeonsworld/ViT-pytorch，它为我们提供了一个很好的起点。

环境要求

pip install timm

vit-explain 快速上手指南

vit-explain 是一个基于 PyTorch 的工具库，用于解释 Vision Transformer (ViT) 模型的注意力机制。它实现了 Attention Rollout 和 Gradient Attention Rollout 方法，帮助开发者可视化模型在图像分类时关注的具体区域。

环境准备

• 系统要求：支持 Python 3.6+ 的操作系统（Linux, macOS, Windows）。
• 前置依赖：
  • Python
  • PyTorch
  • timm (PyTorch Image Models)

安装步骤

1. 确保已安装 PyTorch。如未安装，请访问 PyTorch 官网 获取适合你环境的安装命令。
2. 安装核心依赖库 timm。国内用户推荐使用清华或阿里镜像源加速安装：

pip install timm -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 克隆本仓库代码到本地：

git clone https://github.com/jacobgil/vit-explain.git
cd vit-explain

基本使用

方式一：通过 Python 代码调用

以下示例展示了如何加载预训练的 DeiT-Tiny 模型，并生成针对特定类别（如类别索引 243）的注意力热力图掩码。

from vit_grad_rollout import VITAttentionGradRollout
import torch

# 加载预训练模型
model = torch.hub.load('facebookresearch/deit:main', 
'deit_tiny_patch16_224', pretrained=True)

# 初始化解释器
# discard_ratio: 丢弃低注意力分数的比例 (0-1)
# head_fusion: 注意力头融合策略 ('mean', 'min', 'max')
grad_rollout = VITAttentionGradRollout(model, discard_ratio=0.9, head_fusion='max')

# 生成掩码 (input_tensor 需为预处理后的张量)
mask = grad_rollout(input_tensor, category_index=243)

注意：如果不指定 category_index，将使用普通的 Attention Rollout；指定后则使用针对特定类别的 Gradient Attention Rollout。

方式二：通过命令行运行

你可以直接使用提供的脚本对图片进行分析。

python vit_explain.py --image_path <image path> --head_fusion <mean, min or max> --discard_ratio <number between 0 and 1> --category_index <category_index>

参数说明：

• --image_path: 输入图片的路径。
• --head_fusion: 注意力头融合方式，可选 mean (平均), min (最小), max (最大)。推荐结合 discard_ratio 使用 max。
• --discard_ratio: 噪声过滤比例，范围 0 到 1。数值越高，保留的强注意力区域越集中（例如 0.9 表示只保留前 10% 的注意力）。
• --category_index: 目标类别的索引（ImageNet 类别 ID）。若不填，默认执行普通 Attention Rollout。

示例命令： 分析一张图片中模型对“狗”（类别 243）的关注区域，采用最大融合策略并过滤掉 90% 的低注意力噪声：

python vit_explain.py --image_path examples/dog.jpg --head_fusion max --discard_ratio 0.9 --category_index 243