如何在PyTorch的TransformerEncoderLayer中获取各注意力头分数？

PyTorch原生的nn.TransformerEncoderLayer默认只返回特征输出，不会直接给出注意力分数。要获取每个注意力头的注意力权重，有两种实用方法：

方法一：自定义编码器层，重写forward逻辑 #

继承TransformerEncoderLayer，修改forward方法，让内部的多头注意力层返回并保存注意力权重。

import torch
import torch.nn as nn

# 自定义Transformer编码器层，用于保存注意力分数
class AttentionCapturingEncoderLayer(nn.TransformerEncoderLayer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 存储当前层的注意力分数，形状为 [batch_size, nhead, seq_len, seq_len]
        self.attention_weights = None

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # 调用多头注意力时开启need_weights=True，获取注意力权重
        src2, attn_weights = self.self_attn(src, src, src, attn_mask=src_mask,
                                            key_padding_mask=src_key_padding_mask,
                                            need_weights=True)
        self.attention_weights = attn_weights  # 保存注意力分数
        # 原层的剩余逻辑不变
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

# 定义输入张量：batch_size=2，sequence_length=16，embedding_dim=4
batch_size = 2
sequence_length = 16
input_size = 4  # 需与d_model保持一致
x = torch.randn(batch_size, sequence_length, input_size)

# 配置Transformer参数
d_model = 4  # 匹配输入的embedding维度
nhead = 2  # d_model必须能被nhead整除（4/2=2，每个头维度为2）
d_hid = 64
dropout = 0.1
nlayers = 2

# 使用自定义编码器层构建Transformer
encoder_layers = AttentionCapturingEncoderLayer(d_model=d_model, nhead=nhead, 
                                               dim_feedforward=d_hid, dropout=dropout)
transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers=nlayers)

# 前向传播
output = transformer_encoder(x)

# 提取各层的注意力分数
for idx, layer in enumerate(transformer_encoder.layers):
    attn_weights = layer.attention_weights
    print(f"第{idx+1}层注意力分数形状:", attn_weights.shape)
    # 可通过attn_weights[b, h]获取第b个样本、第h个注意力头的注意力矩阵

方法二：用钩子函数捕获注意力分数 #

无需修改原类，通过给多头注意力层注册前向钩子，动态捕获注意力权重，更灵活。

import torch
import torch.nn as nn

# 定义输入张量
batch_size = 2
sequence_length = 16
input_size = 4
x = torch.randn(batch_size, sequence_length, input_size)

# 配置Transformer参数
d_model = 4
nhead = 2
d_hid = 64
dropout = 0.1
nlayers = 2

encoder_layers = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, 
                                           dim_feedforward=d_hid, dropout=dropout)
transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers=nlayers)

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

# 定义钩子函数，捕获注意力权重
def save_attention(module, input, output):
    # 强制开启返回注意力权重
    module.need_weights = True
    attn_weights = output[1]
    # 定位当前模块所属的层
    for idx, layer in enumerate(transformer_encoder.layers):
        if layer.self_attn is module:
            attention_weights_dict[idx] = attn_weights
            break

# 给每一层的多头注意力模块注册钩子
for layer in transformer_encoder.layers:
    layer.self_attn.register_forward_hook(save_attention)

# 前向传播
output = transformer_encoder(x)

# 输出各层注意力分数信息
for layer_idx, attn_weights in attention_weights_dict.items():
    print(f"第{layer_idx+1}层注意力分数形状:", attn_weights.shape)

注意事项 #

• d_model必须与输入的embedding维度（这里是4）完全一致，且nhead必须能整除d_model，否则会抛出维度不匹配的错误。
• 注意力分数的形状为[batch_size, nhead, sequence_length, sequence_length]，其中attn_weights[b, h, i, j]代表第b个样本中，第h个注意力头里，第i个token对第j个token的注意力权重。

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