如何在PyTorch中获取预训练Transformer的N×N注意力分数矩阵？

以下是两种仅用PyTorch实现的方法，无需依赖BertViz等外部工具：

方法一：针对Hugging Face预训练模型（如BERT）使用Forward Hook捕获 #

Hugging Face的Transformer注意力层（如BertSelfAttention）的forward方法实际会返回注意力输出和softmax后的注意力分数矩阵，只是上层模块默认只取用前者。我们可以通过重写注意力层的forward并注册hook来保存这个矩阵：

import torch
from transformers import BertModel, BertTokenizer

# 初始化模型与tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
model.eval()

# 用于存储各层注意力分数的字典
attention_store = {}

def capture_attention(module, _, __):
    attention_store[f"layer_{module.layer_idx}"] = module.saved_attn_probs

# 遍历所有encoder层，修改注意力模块并注册hook
for layer_idx, layer in enumerate(model.encoder.layer):
    attn_module = layer.attention.self
    attn_module.layer_idx = layer_idx  # 标记层索引

    # 重写forward方法，保存注意力分数
    original_forward = attn_module.forward
    def new_forward(self, hidden_states, attention_mask=None, head_mask=None):
        # 调用原forward，获取完整输出（attention_output, attention_probs）
        full_output = original_forward(hidden_states, attention_mask, head_mask)
        self.saved_attn_probs = full_output[1].detach().cpu()
        return full_output[0]  # 返回原上层需要的输出

    # 绑定新方法到当前注意力模块
    attn_module.forward = new_forward.__get__(attn_module, type(attn_module))
    # 注册hook
    attn_module.register_forward_hook(capture_attention)

# 测试输入
text = "Hello world"
inputs = tokenizer(text, return_tensors="pt")

with torch.no_grad():
    _ = model(**inputs)

# 查看结果：每个层的注意力分数形状为 (batch_size, num_heads, seq_len, seq_len)
print(attention_store["layer_0"].shape)  # 示例输出: torch.Size([1, 12, 3, 3])

方法二：重写注意力层类并替换模型组件 #

通过继承原注意力层类，在forward中直接保存注意力分数，再替换模型中的对应层：

import torch
from transformers import BertModel, BertTokenizer
from transformers.models.bert.modeling_bert import BertSelfAttention

class AttnCapturingBertSelfAttention(BertSelfAttention):
    def __init__(self, config):
        super().__init__(config)
        self.attention_probs = None

    def forward(self, hidden_states, attention_mask=None, head_mask=None):
        full_output = super().forward(hidden_states, attention_mask, head_mask)
        self.attention_probs = full_output[1].detach().cpu()
        return full_output[0]

# 初始化模型并替换注意力层
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
model.eval()

for layer in model.encoder.layer:
    layer.attention.self = AttnCapturingBertSelfAttention(model.config)

# 测试输入
text = "Hello world"
inputs = tokenizer(text, return_tensors="pt")

with torch.no_grad():
    _ = model(**inputs)

# 获取各层注意力分数
for idx, layer in enumerate(model.encoder.layer):
    print(f"Layer {idx} 注意力矩阵形状: {layer.attention.self.attention_probs.shape}")

针对原生PyTorch nn.Transformer的情况 #

原生nn.MultiheadAttention支持通过need_weights=True参数直接返回注意力分数矩阵：

import torch
import torch.nn as nn

# 初始化原生Transformer
transformer = nn.Transformer(d_model=512, nhead=8)
src = torch.rand(10, 32, 512)  # 输入序列: (seq_len, batch_size, d_model)
tgt = torch.rand(20, 32, 512)  # 目标序列

with torch.no_grad():
    output, attn_weights = transformer(src, tgt, need_weights=True)

# 注意力权重形状: (num_heads * batch_size, tgt_len, src_len)
print(attn_weights.shape)  # 示例输出: torch.Size([256, 20, 10])

关键说明 #

捕获到的注意力矩阵通常为(batch_size, num_heads, seq_len, seq_len)格式，其中每个注意力头对应一个N×N的分数矩阵（N为输入序列长度）。

内容的提问来源于stack exchange，提问作者Thomas