10.7 从零实现BERT#
经过10.6节内容的介绍，我们对于BERT模型的整体结构已经有了一定的了解。根据图10-26可知，从本质上来说BERT就是由Transformer中的编码器构建而来，同时在输入层部分额外加入了一个句子编码来区分输入的不同部分。在本节内容中，我们将以图10-26中黑色加粗字体所示的部分为一个类来分别构建实现整个BERT模型。
10.7.1 工程结构#
由于整个项目涉及到的代码模块较多，所以我们在这里先进行简单说明， 这样便于各位读者在阅读后续内容时能够快速地定位到相应的代码部分。同时， 这里也建议各位读者在阅读内容的能够同时结合代码一起阅读并动手实践。
在整个工程项目中一共包含有6个主要的文件目录，cache、data、model、Tasks、test和uils。其中cache目录用来存放训练过程中所保存下来的模型；data用来存放各类数据集，包括后续将要用到的文本分类、问题回答和训练预料等；model目录中存放的是整个BERT模型的实现代码，以及相关下游任务的构造模型；Tasks目录中存放的是model中各个任务对应的模型训练代码；test目录中存放的是各个模块的测试案例，用于在实现过程中的验证，后续内容中的相关使用示例都能在其中找到；utils目录中存放的是一些辅助工具模块，包括数据集构建和日志模块等。
10.7.2 Input Embedding实现#
首先，我们先来看Input Embedding的实现过程。由于在10.4节内容中我们已经介绍过了字符嵌入的实现，所以在复用这部分代码之后只需要再实现位置编码和句子编码即可。本节内容及后续多个下游任务的完整示例代码可参见Code/Chapter10/C04_BERT文件。
1. Positional Embedding实现
不同于Transformer中位置编码的实现方式，在BERT中位置编码并没有采用固定的变换公式来计算每个位置上的值，而是采用了普通嵌入层的方式来为每个位置生成一个向量然后随着模型一起训练。因此，这也就限制了在使用预训练的中文BERT模型时最大的序列长度只能是512，因为在训练时只初始化了512个位 置向量。示例代码如下所示：
1 class PositionalEmbedding(nn.Module):
2 def __init__(self, hidden_size, max_position_embeddings=512,
3 initializer_range=0.02):
4 super(PositionalEmbedding, self).__init__()
5 self.embedding = nn.Embedding(max_position_embeddings, hidden_size)
6
7 def forward(self, position_ids):
8 return self.embedding(position_ids).transpose(0, 1)在上述代码中，第7行position_ids的形状为[1,position_ids_len]。第8行返回结果的形状为[position_ids_len, 1, hidden_size]。
2. Segment Embedding 实现
句子编码是对输入的两个序列分别赋予一个位置向量用以区分各自所在的位置，这一点可以和上面的位置编码进行类比。具体地，示例代码如下所示：
1 class SegmentEmbedding(nn.Module):
2 def __init__(self, type_vocab_size, hidden_size, initializer_range=0.02):
3 super(SegmentEmbedding, self).__init__()
4 self.embedding = nn.Embedding(type_vocab_size, hidden_size)
5
6 def forward(self, token_type_ids):
7 return self.embedding(token_type_ids)在上述代码中，第2行type_vocab_size的默认值为 2，即只用于区分两个序列的不同位置。第6行token_type_ids的形状为[token_type_ids_len, batch_size]。第7行返回结果的形状为[token_type_ids_len, batch_size, hidden_size]。
3. Bert Embedding实现
在完成3个部分的代码实现之后，只需要将每个部分的结果相加便可以得到最终的嵌入层表示作为BERT模型的输入，示例代码如如下所示：
1 class BertEmbeddings(nn.Module):
2 def __init__(self, config):
3 super().__init__()
4 self.word_embeddings = TokenEmbedding(config.vocab_size,
5 config.hidden_size,config.pad_token_id)
6 self.position_embeddings = PositionalEmbedding(
7 config.max_position_embeddings,config.hidden_size)
8 self.token_type_embeddings = SegmentEmbedding(config.type_vocab_size,
9 config.hidden_size,config.initializer_range)
10 self.LayerNorm = nn.LayerNorm(config.hidden_size)
11 self.dropout = nn.Dropout(config.hidden_dropout_prob)
12 self.register_buffer("position_ids",
13 torch.arange(config.max_position_embeddings).expand((1, -1)))在上述代码中，第2行config是传入的一个配置类，里面各个类成员就是BERT模型中对应的模型参数。第 4~8行是分别用来定义图10-27中的3个编码部分。第12~13行是用来生成一个默认的位置编号，即[0,1,....,511]。
进一步，其前向传播过程代码为：
1 def forward(self,input_ids=None,position_ids=None,token_type_ids=None):
2 src_len = input_ids.size(0)
3 token_emb = self.word_embeddings(input_ids)
4 if position_ids is None:
5 position_ids = self.position_ids[:, :src_len]
6 positional_emb = self.position_embeddings(position_ids)
7 if token_type_ids is None:
8 token_type_ids = torch.zeros_like(input_ids)
9 segment_emb = self.token_type_embeddings(token_type_ids)
10 embeddings = token_emb + positional_emb + segment_emb
11 embeddings = self.LayerNorm(embeddings)
12 embeddings = self.dropout(embeddings)
13 return embeddings在上述代码中，第1行input_ids表示输入序列的原始索引编号，即根据词表映射后的索引形状为[src_len, batch_size]。第4~6行position_ids是位置序列， 本质是[0,1,2,3,...,src_len-1]形状为[1,src_len]，在实际建模时这个参数可以不用传值，因为当其为空时会自动从self.position_ids截取一段。第7~9行token_type_ids用于不同序列之间的分割，例如[0,0,0,0,1,1,1,1]用于区分前后不同的两个句子形状为[src_len,batch_size]。如果输入模型的只有一个序列，那么这个参数也不用传值。第 10~12行代码则是用来将3部分的编码结果进行相加。
4. 使用示例
在实现完上述代码之后，便可以通过如下方式进行使用：
1 if __name__ == '__main__':
2 json_file = '../bert_base_chinese/config.json'
3 config = BertConfig.from_json_file(json_file)
4 src = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]], dtype=torch.long)
5 token_type_ids = torch.LongTensor([[0, 0, 0, 1, 1], [0, 0, 1, 1, 1]])
6 src, token_type_ids = src.transpose(0, 1), token_type_ids.transpose(0, 1)
7 bert_embedding = BertEmbeddings(config)
8 bert_embedding_result = bert_embedding(src, token_type_ids=token_type_ids)
9 print(bert_embedding_result.shape) # torch.Size([5, 2, 768])在上述代码中，第2~3行是载入原始的BERT模型配置文件，里面包含了hidden_size、max_position_embeddings等默认参数的取值。第4~6行是生成输入层对应的输入部分。第7~9行是实例化BERT嵌入层并计算前向传播的输出结果，形状为[src_len, batch_size, hidden_size]。
10.7.3 BERT网络实现#
在实现完Input Embedding部分的代码后，下面可以着手来实现构成BERT模型的第2个重要组成部分BertEncoder。如图10-26所示，整个 BertEncoder由多个BertLayer堆叠形成；而BertLayer又是由BertOutput、BertIntermediate和BertAttention这3个部分组成；而BertAttention是由BertSelfAttention和BertSelfOutput所构成。之所以需要将整个模型拆分成各个模块进行实现，主要是为了降低功能模块之间的耦合性，以便按需进行调整。
1. BertaAttention实现
对于BertAttention来说其核心是Transformer中所提出的self-attention机制，即图10-26中的BertSelfAttention模块；其次是一个残差连接和标准化操作。对于BertSelfAttention的实现，示例代码如下所示：
1 class BertSelfAttention(nn.Module):
2 def __init__(self, config):
3 super(BertSelfAttention, self).__init__()
4 if 'use_torch_multi_head' in config.__dict__ and config.use_torch_multi_head:
5 MultiHeadAttention = nn.MultiheadAttention
6 else:
7 MultiHeadAttention = MyMultiheadAttention
8 self.multi_head_attention = MultiHeadAttention(config.hidden_size,
9 config.num_attention_heads,config.attention_probs_dropout_prob)
10
11 def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
12 return self.multi_head_attention(query, key, value,
13 attn_mask=attn_mask, key_padding_mask=key_padding_mask)在上述代码中，第4~10行是实例化一个多头注意力机制对象，并且这里我们提供了两种多头实现，一种10.3节内容中介绍的都头注意力实现，另一种是直接使用PyTorch框架中的默认实现，可以通过设置参数use_torch_multi_head = True进行切换。第12~15行则是多头注意力的前向传播过程，其返回包含两个部分，多头注意力的线性组合以及多头注意力权重的均值，形状分别为[tgt_len, batch_size, hidden_size]和[batch_size, tgt_len, src_len]。
进一步，对于BertSelfOutput的实现包括层 Dropout、标准化和残差连接3个操作，示例代码如下：
1 class BertSelfOutput(nn.Module):
2 def __init__(self, config):
3 super().__init__()
4 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)
5 self.dropout = nn.Dropout(config.hidden_dropout_prob)
6
7 def forward(self, hidden_states, input_tensor):
8 hidden_states = self.dropout(hidden_states)
9 hidden_states = self.LayerNorm(hidden_states + input_tensor)
10 return hidden_states上述代码便是BertSelfOutput的实现，其过程也十分简单这里就不再赘述，最后第10行返回结果的形状为[src_len, batch_size, hidden_size]。
接下来就是对BertAttention部分进行实现，其由BertSelfAttention和BertSelfOutput这两个类构成，示例代码如下所示：
1 class BertAttention(nn.Module):
2 def __init__(self, config):
3 super().__init__()
4 self.self = BertSelfAttention(config)
5 self.output = BertSelfOutput(config)
6
7 def forward(self,hidden_states,attention_mask=None):
8 self_outputs = self.self(hidden_states,hidden_states,hidden_states,
9 attn_mask=None,key_padding_mask=attention_mask)
10 attention_output = self.output(self_outputs[0], hidden_states)
11 return attention_output在上述代码中，第7行hidden_states是输入层处理后的结果，形状为[src_len, batch_size, hidden_size]，attention_mask是同一个小批量样本中不同长度序列的掩码填充信息，即在10.4节内容中所介绍的key_padding_mask，形状为[batch_size, src_len]，这里只是为了和PyTorch中的命名方式保持一致。第8~9行是自注意力机制的输出结果。第10~11行便是执行BertSelfOutput中的3个操作，最后返回结果的形状为[src_len, batch_size, hidden_size]。
2. BertLayer实现