快速获取hugging face的LLM推理时的中间层attention weights的办法

以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库的模型自带的功能，用起来很方便。