bert模型取last_hidden_state[:, 0]

原创已于 2023-01-05 16:51:39 修改 · 2.6k 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#bert #python #深度学习

于 2023-01-05 16:39:58 首次发布

文章解释了在处理模型输出时，如何从(32,100,768)的形状转换到(32,768)，涉及到batch_size、max_length和hidden_size的概念。通过一个numpy数组的例子，展示了last_hidden_state[:,0]如何提取每个样本的第一行，即(1,768)形状的数据，最终得到(32,768)的形状，这通常与BERT模型中的CLStoken相关联。

问题

看代码时,第一行跑完target_pred的last_hidden_state的shape为(32,100,768)，第二行跑完target_pred的shape为(32,768) ，不理解。
注：32为batch_size，100为max_length，768为hidden_size。

target_pred = model(target_input_ids, target_attention_mask, target_token_type_ids, output_hidden_states=True, return_dict=True)
target_pred = target_pred.last_hidden_state[:, 0]

解决

涉及的知识点是三维数组切片，测试代码如下：

a=np.arange(0, 24, 1).reshape(3,2,4)
print(a)
# 输出
[[[ 0  1  2  3]
  [ 4  5  6  7]]

 [[ 8  9 10 11]
  [12 13 14 15]]

 [[16 17 18 19]
  [20 21 22 23]]]
  
b=a[:,0]
print(b)
print(b.shape)
# 输出
[[ 0  1  2  3]
 [ 8  9 10 11]
 [16 17 18 19]]
(3, 4)

以上面的例子类推，last_hidden_state[:, 0]的含义：每个batch有32个样本，每个样本的shape为(100,768)的二维数组，对每个样本取第一行，即(1,768)，因此last_hidden_state[:, 0]的shape为(32,768)，相当于降维。并且调用了cls。