如何从BERT模型获取隐藏层/状态输出并转为数组？

问题 #

我通过以下代码设置output_hidden_states=True并获取了BERT模型的输出：

from transformers import BertModel, BertConfig

config = BertConfig.from_pretrained("xxx", output_hidden_states=True)
model = BertModel.from_pretrained("xxx", config=config)

outputs = model(inputs)

输出是BaseModelOutputWithPoolingAndNoAttention类的实例，包含last_hidden_state、pooler_output、hidden_states三个属性。我不确定last_hidden_state的类型，想知道怎么获取它的值并转为类似[0,1,2,3,...]的普通数组？

解答 #

• last_hidden_state的类型：它是PyTorch的Tensor对象（若使用TensorFlow框架则为tf.Tensor），形状通常为[batch_size, sequence_length, hidden_size]，对应批量数据中每个序列的每个token的隐藏层输出结果。
• 转换为普通数组的方法：
  1. 若为PyTorch张量：先调用.detach()脱离计算图（训练模式下需要这一步避免梯度关联），再用.numpy()转为NumPy数组，最后通过.tolist()得到Python列表：

     last_hidden_list = outputs.last_hidden_state.detach().numpy().tolist()
     

  2. 若为TensorFlow张量：直接用.numpy()转为NumPy数组，再转成Python列表即可：

     last_hidden_list = outputs.last_hidden_state.numpy().tolist()
     

• GPU运行时的额外处理：如果模型在GPU上运行，需要先将张量移至CPU再转换，PyTorch下可添加.cpu()：

  last_hidden_list = outputs.last_hidden_state.cpu().detach().numpy().tolist()
  

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