BERT模型结构可视化与模块维度转换剖析_bert输出维度

  大家好，我是herosunly。985院校硕士毕业，现担任算法研究员一职，热衷于机器学习算法研究与应用。曾获得阿里云天池比赛第一名，科大讯飞比赛第三名，CCF比赛第四名。拥有多项发明专利。对机器学习和深度学习拥有自己独到的见解。曾经辅导过若干个非计算机专业的学生进入到算法行业就业。希望和大家一起成长进步。

  本文作为BERT模型讲解的第一篇文章，主要介绍了BERT模型结构可视化显示与模块剖析，希望对新手有所帮助。

1. 可视化显示

1.1 安装依赖库

  首先需要说明的是要提前安装好PyTorch(版本大于等于1.7)，其他依赖库安装方法如下：

pip install graphviz
pip install torchview
1
2

  Successfully installed torchview则说明torchview安装成功了，如下图所示：

1.2 可视化

  由于BERT-Base模型层数较深，这里以BERT-Tiny为例进行展示，其中BERT-Tiny的模型层数为2(Transformer blocks)、自注意力头的个数为2、隐藏层节点数为128(hidden-size)，BERT-Base对应的三大参数如BERT论文所示：

  首先设置模型缓存路径，代码如下所示：

import os
os.environ['TRANSFORMERS_CACHE'] = '/home/model/transformers/cache'
1
2

  然后设置可视化结果为png图片，如下所示：

import graphviz
graphviz.set_jupyter_format('png')
1
2

  最后进行模型可视化，其中输入序列为三句话，即Hello new world、good good study和ha ha ha。

from transformers import AutoModel, AutoTokenizer
from torchview import draw_graph

tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")
model_tiny = AutoModel.from_pretrained("prajjwal1/bert-tiny")

inputs = tokenizer(["hello new world", 'good good study', 'ha ha ha'], return_tensors="pt")

model_graph = draw_graph(model_tiny, input_data=inputs)

model_graph.visual_graph
1
2
3
4
5
6
7
8
9
10
11

2. 模块维度转换剖析

2.1 输入文本维度

  首先input-tensor维度为(3, 5)，其中3指代的是句子的个数(3)，5指代的是token的个数，那么为什么是5呢？

  由于在tokenizer中是存在一些特殊的tokens的，如[UNK]、[SEP]、[PAD]、[CLS]、[MASK]，可通过下列代码进行显示，如下所示：

print(tokenizer.special_tokens_map)
1

  它们对应的id为多少呢，代码如下所示：

tokenizer.convert_tokens_to_ids(['[UNK]', '[SEP]', '[PAD]', '[CLS]', '[MASK]'])
1

[100, 102, 0, 101, 103]
1

  然后我们再看看看输入文本对应的维度(3, 5)，可以看到每个序列最前面的token id为101，最后面的token id为102,，如下图所示：

  对应之前介绍的特殊tokens，可得对原有文本基础上在开头添加了[CLS]，而在末尾添加了[SEP]。所以文本长度从3增加到了5。

2.2 Embedding模块

  Embedding模块本质上是由Token Embeddings + Segment Embeddings+Position Embeddings，其中每一种Embeddings的维度大小均和hidden size是一致的，如下图所示：

  下图即为Embedding的运算图，图中第二层中的Embedding分别为Token Embeddings和Position Embeddings，由于每一种Embeddings的维度大小均为128(hidden size)，所以两者相加后维度为(3, 5, 128)。加上Segment Embeddings后维度不变，具体如下图所示：

2.3 LayerNorm

  由于LayerNorm和dropout并不改变原有数据维度，所以则和之前Embedding模块的输出维度是一致的，即(3, 5, 128)。

2.4 Self-attention模块

2.4.1 Attention Mask模块

  在第一幅总图中的右侧中，又对输入文本进行了处理，但很多同学对此部分并不了解，本质上它是Attention Mask模块。具体将在下文中进行介绍。

2.4.1.1 关键概念之序列填充(padding)

  在神经网络的训练中，为了加速训练并且确保网络能够达到收敛的状态，往往是进行批量样本(batch)的学习，也就是说输入样本通过张量(矩阵)的形式进行输入，则张量的维度需要保持一致，最终需要使得必须让同一个batch中的每个样本的序列长度一致。

  在保持序列长度一致性的方法中，最常见的就是通过padding补零，把同一个batch中的所有样本都变成同一个长度，从而便于批量计算。对于填充值(如零)，可使用mask机制来避免模型对填充值进行训练。

  那么问题来了，既然是填充零值，是往哪个方向填充呢？是往后面进行填充(post)呢还是往前面进行填充(pre)。根据huggingface的BERT文档中第一页的前半部分，链接为https://huggingface.co/docs/transformers/model_doc/bert，如下图所示：

  将原有输入句子中的第三句话中的最后一个单词ha进行删除，如下所示：

new_inputs = tokenizer(["hello new world", 'good good study', 'ha ha'], return_tensors="pt", padding=True)
print(new_inputs)
1
2

  可以看到attention_mask第三句话中的最后一位为0，即上述讲解中的序列填充(padding)的概念。之所以对填充位赋值为0，简单来说，也就是填充位在self-attention模块中是不需要进行计算的。

2.4.2 Self-attention模块

$$Attention(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$


  https://github.com/google-research/bert中的modeling.py中的self-attention代码如下：

def transpose_for_scores(input_tensor, batch_size, num_attention_heads,
                       seq_length, width):
    output_tensor = tf.reshape(
        input_tensor, [batch_size, seq_length, num_attention_heads, width])

    output_tensor = tf.transpose(output_tensor, [0, 2, 1, 3])
    return output_tensor
1
2
3
4
5
6
7

  可以看到Q对应的维度为[batch_size, num_attention_heads, seq_length, width]，在本篇文本中对应的维度则为[3, 2, 5, 64]。其中batch_size=3，这里可以认为和句子个数(3)是相同的。

  同理K和V对应的维度亦为[3, 2, 5, 64]，$K^T$对应的维度为[3, 2, 64, 5]，则$Q{K}^T$对应的维度为[3, 2, 5, 5] (这里类似于PyTorch中的bmm操作)。在softmax之前会先和attention_mask进行相加。

if attention_mask is not None:
  attention_mask = tf.expand_dims(attention_mask, axis=[1])

  adder = (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0
  attention_scores += adder
1
2
3
4
5

  原始的attention_mask的维度是[3, 5]，如下图所示：

  经过维度扩充后变为[3, 1, 1, 5]，即图中的__getitem__操作，to指代的是tf.cast(attention_mask, tf.float32)，然后__rsub__对应上述代码中的1.0 - xxx，mul对应的是* -10000.0操作，那么问题来了，为什么要将维度转换为[3, 1, 1, 5]呢？本质上是为了和[3, 2, 5, 5]进行add操作，此时固定的是维度中的3和5，然后对1和1进行广播操作，从而得到[3, 2, 5, 5]。

  softmax和算术操作不改变数据维度，则$Softmax(\frac{Q{K}^T}{\sqrt{d_k}}）$维度为[3, 2, 5, 5]，再乘以$V$对应的维度为[3, 2, 5, 64]，通过维度转置操作后为[3, 5, 2, 64]，再通过维度合并后为[3, 5, 128]。

query_layer = transpose_for_scores(query_layer, batch_size,
                                 num_attention_heads, seq_length,
                                 size_per_head)

key_layer = transpose_for_scores(key_layer, batch_size, num_attention_heads,
                               seq_length, size_per_head)


attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
attention_scores = tf.multiply(attention_scores,
                             1.0 / math.sqrt(float(size_per_head)))

if attention_mask is not None:
  attention_mask = tf.expand_dims(attention_mask, axis=[1])

  adder = (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0
  attention_scores += adder

attention_probs = tf.nn.softmax(attention_scores)
value_layer = tf.reshape(
    value_layer,
    [batch_size, seq_length, num_attention_heads, size_per_head])

value_layer = tf.transpose(value_layer, [0, 2, 1, 3])

context_layer = tf.matmul(attention_probs, value_layer)
context_layer = tf.transpose(context_layer, [0, 2, 1, 3])

context_layer = tf.reshape(
        context_layer,
        [batch_size, seq_length, num_attention_heads * size_per_head]])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

2.5 Position-wise Feed-Forward Networks

  这里的$d_{ff}$即为512，即$W_1$对应为512(BertIntermediate的维度)，$W_2$对应维度为128(BertOutput的维度)，如下图所示：

2.6 最终输出

  模型默认输出如右边所示是(3, 5, 128)，即黄色部分，左边对整个句子所有token的向量进行了平均表示，所以维度变为了(3, 128)。