以GPT-2为例,只需要在推理时加上output_attentions = True
即可:
from transformers import GPT2Model, GPT2Config
configuration = GPT2Config(
n_positions= 200,
n_embd=256,
n_layer=12,
n_head=8,
resid_pdrop=0.0,
embd_pdrop=0.0,
attn_pdrop=0.0,
use_cache=False,
) # GPT-2配置
model = GPT2Model(configuration)
input_embedding = torch.randn(64, 200, 256) # [bsz, sequence_len, n_embd]
pred = model(inputs_embeds=embeds, output_attentions = True) #前传加上output_attentions = True即可
加上output_attentions = True,则返回的pred为包含logits, attentions等key的字典。其中的 attentions 是一个长度为模型层数的元组,其中第i个元素是一个形状为[bsz, n_head, sequence_len, sequence_len]的tensor,即为第i层的attention weight。这是transformer库的模型自带的功能,用起来很方便。