Bert源码阅读

前言

对Google开源出来的bert代码，来阅读下。不纠结于代码组织形式，而只是梳理下其训练集的生成，训练的Transformer和multi-head实现，它的具体实现和论文里写的还是有很大差别的。

训练集的生成

主要实现在create_pretraining_data.py和tokenization.py两个脚本里。
输入文本格式举例，下面是两篇文章外加一篇空文章。两篇文章之间用空格作间隔。

This is a blog about bert code reading.
It is writed using markdown, which is a markup language that can be written using a plain text editor.
Hopefuly it will give the reader a deep understanding of bert.

本文是篇关于bert源码阅读的博客。
它是用markdown写的，markdown是种可以使用普通文本编辑器编写的标记语言。
希望本文能够给读者以对bert更深层次的理解。

第一步，读取raw文本，按行分词处理后存储all_documents[doc_0, doc_1, …]里面，doc_i=[line_0, line_1, …], line_i = [token_0, token_1, …]，然后shuffle文章。
第二步，重复dupe_factor=10次，每篇文章生成样本，[CLS +A+SEP +B+SEP]作一条样本。

  for _ in range(dupe_factor):
    for document_index in range(len(all_documents)):
      instances.extend(
          create_instances_from_document(
              all_documents, document_index, max_seq_length, short_seq_prob,
              masked_lm_prob, max_predictions_per_seq, vocab_words, rng))

create_instances_from_document函数对每篇文章都生成一个训练样本实例。
从第一条句子循环到最后一条句子$i$，收集segment到current_chunk列表中，当收集到的总句子长度>=单条样本最长值时，构造A+B。

 if i == len(document) - 1 or current_length >= target_seq_length:

随机截取 current_chunk的某个位置a_end，[0, a_end]作为子句A=token_a。
B句随机概率选择是Next or Not next，如果是next，则current_chunk的剩余[a_end, :]作为子句B=token_b。如果Not next，则随机挑一篇文章，选择某个长度的子句作为B=token_b。注意，Not next时，循环经过的B句子对应的步幅，要回去(因为这部分句子并没有被真正使用，所以退回去以免浪费)。

          num_unused_segments = len(current_chunk) - a_end
          i -= num_unused_segments

两个句子加和长度超过最大长度怎么办？使用truncate_seq_pair在A和B中随机选择一个，随机丢掉首/尾的词，每次丢一个token，直到加和长度<=最大长度。

truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng)

之后根据token_a和token_b生成tokens和segment_ids
$tokens = [CLS, A_0, A_1, A_2, SEP, B_0, B_1, B_2, SEP]$
$segment\_ids =[0_a, 0_a, 0_a, 0_a, 0_a, 1_b, 1_b, 1_b, 1_b]$
再之后，根据tokens生成遮挡之后的tokens、遮挡位置masked_lm_positions、遮挡位置的真实词masked_lm_labels。

        (tokens, masked_lm_positions,
         masked_lm_labels) = create_masked_lm_predictions( 
         tokens, masked_lm_prob, max_predictions_per_seq, vocab_words, rng)

15%采样遮挡，对遮挡的处理情况如下：
a) 80%的概率，遮挡词被替换为[mask]。$\longrightarrow$别人看不到我。
b) 10%的概率，遮挡词被替换为随机词。$\longrightarrow$别人看走眼我。
c) 10%的概率，遮挡词被替换为原来词。$\longrightarrow$别人能看到我。

    masked_token = None
    # 80% of the time, replace with [MASK]
    if rng.random() < 0.8:
      masked_token = "[MASK]"
    else:
      # 10% of the time, keep original
      if rng.random() < 0.5:
        masked_token = tokens[index]
      # 10% of the time, replace with random word
      else:
        masked_token = vocab_words[rng.randint(0, len(vocab_words) - 1)]

输入和返回结果举例：

input tokens  ="The man went to the store . He bought a gallon of milk "
ouput tokens ="The man went to the [mask] . He [mask] a gallon of milk"
output masked_lm_positions = [5, 8, 10, 11]
output masked_lm_labels = [store, bought, gallon, ice]

位置#5,#8被遮挡，#10被替换为原token，#11被替换为随机词。注意CLS和SEP不会被遮挡。
然后保存成TrainingInstance类，同时保留了is_next标记.

        instance = TrainingInstance(
            tokens=tokens,
            segment_ids=segment_ids,
            is_random_next=is_random_next,
            masked_lm_positions=masked_lm_positions,
            masked_lm_labels=masked_lm_labels)

tokenization.FullTokenizer类用来处理分词，标点符号，unknown词，Unicode转换等操作。注意：中文只有单个字的切分，没有词。

数据存储及读取

存储为TF-Record
输入sentence变量的处理

input_ids = tokenizer.convert_tokens_to_ids(instance.tokens)  ## ID化 ##
input_mask = [1] * len(input_ids)
segment_ids = segment_ids
padding 0 --> max_seq_length
1. 对iput_ids 补0到句子最大长度
2. 对input_mask 补0到句子最大长度
3. 对segment_ids 补0到句子最大长度

注意：input_mask是样本中有效词句的标识，后面需要用作作attention视野的约束。
遮挡变量的处理

    masked_lm_positions = list(instance.masked_lm_positions)
    masked_lm_ids = tokenizer.convert_tokens_to_ids(instance.masked_lm_labels)
    masked_lm_weights = [1.0] * len(masked_lm_ids)
    ## padding 0 --> max_seq_length

注意：masked_lm_ids是有mask的词对应的ID；masked_lm_positions是有mask的词对应的句子中位置。
next_sentense 处理

next_sentence_label = 1 if instance.is_random_next else 0

save format 处理

    features = collections.OrderedDict()
    features["input_ids"] = create_int_feature(input_ids)
    features["input_mask"] = create_int_feature(input_mask)
    features["segment_ids"] = create_int_feature(segment_ids)
    features["masked_lm_positions"] = create_int_feature(masked_lm_positions)
    features["masked_lm_ids"] = create_int_feature(masked_lm_ids)
    features["masked_lm_weights"] = create_float_feature(masked_lm_weights)
    features["next_sentence_labels"] = create_int_feature([next_sentence_label])

    tf_example = tf.train.Example(features=tf.train.Features(feature=features))

读取使用dataset。

    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]
    masked_lm_positions = features["masked_lm_positions"]
    masked_lm_ids = features["masked_lm_ids"]
    masked_lm_weights = features["masked_lm_weights"]
    next_sentence_labels = features["next_sentence_labels"]

BertModel

模型实例化 ，注意这里的变量对应。

    model = modeling.BertModel(
        config=bert_config,
        is_training=is_training,
        input_ids=input_ids,
        input_mask=input_mask,
        token_type_ids=segment_ids, ## token_type是句子标记 ##
        use_one_hot_embeddings=use_one_hot_embeddings)

输入token_ids–>向量化处理, embeding_lookup返回token_emb 和查询的table表。

        (self.embedding_output, self.embedding_table) = embedding_lookup(
            input_ids=input_ids,
            vocab_size=config.vocab_size,
            embedding_size=config.hidden_size,
            initializer_range=config.initializer_range,
            word_embedding_name="word_embeddings",
            use_one_hot_embeddings=use_one_hot_embeddings)

加入pos_emb和type_emb处理， embedding_postprocessor
注意：pos_emb并不是用sin/cos函数生成的，而是随机生成的。

        self.embedding_output = embedding_postprocessor(
            input_tensor=self.embedding_output,
            use_token_type=True, ## type_emb的处理设置 ##
            token_type_ids=token_type_ids,
            token_type_vocab_size=config.type_vocab_size,
            token_type_embedding_name="token_type_embeddings",
            use_position_embeddings=True, ## pos_emb的处理设置 ##
            position_embedding_name="position_embeddings",
            initializer_range=config.initializer_range,
            max_position_embeddings=config.max_position_embeddings,
            dropout_prob=config.hidden_dropout_prob)

重要：构造attention可视域的attention_mask，因为每个样本都经过padding了，视野必须要约束到有效范围词句内。

        # This converts a 2D mask of shape [batch_size, seq_length] to a 3D
        # mask of shape [batch_size, seq_length, seq_length] which is used
        # for the attention scores.
        attention_mask = create_attention_mask_from_input_mask(input_ids, input_mask)
        ## 注意： 
        ## input_ids 是经过padding后的 [32,108, 99, 0, 0]; ##
        ## input_mask 是有效词标志     [1, 1,   1,  0, 0] ##
def create_attention_mask_from_input_mask(from_tensor, to_mask):
  """Create 3D attention mask from a 2D tensor mask.
  Args:
    from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...].
    to_mask: int32 Tensor of shape [batch_size, to_seq_length].
  Returns:
    float Tensor of shape [batch_size, from_seq_length, to_seq_length].
  """
  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3]) 
  batch_size = from_shape[0]
  from_seq_length = from_shape[1]

  to_shape = get_shape_list(to_mask, expected_rank=2)
  to_seq_length = to_shape[1]

  to_mask = tf.cast(
      tf.reshape(to_mask, [batch_size, 1, to_seq_length]), tf.float32)
  # We don't assume that `from_tensor` is a mask (although it could be). We
  # don't actually care if we attend *from* padding tokens (only *to* padding)
  # tokens so we create a tensor of all ones.
  # `broadcast_ones` = [batch_size, from_seq_length, 1]
  broadcast_ones = tf.ones(
      shape=[batch_size, from_seq_length, 1], dtype=tf.float32)
  # Here we broadcast along two dimensions to create the mask.
  mask = broadcast_ones * to_mask

  return mask

Bert.Transformer

        # Run the stacked transformer.
        # `sequence_output` shape = [batch_size, seq_length, hidden_size].
        self.all_encoder_layers = transformer_model(
            input_tensor=self.embedding_output,
            attention_mask=attention_mask,
            hidden_size=config.hidden_size,
            num_hidden_layers=config.num_hidden_layers,
            num_attention_heads=config.num_attention_heads,
            intermediate_size=config.intermediate_size,
            intermediate_act_fn=get_activation(config.hidden_act),
            hidden_dropout_prob=config.hidden_dropout_prob,
            attention_probs_dropout_prob=config.attention_probs_dropout_prob,
            initializer_range=config.initializer_range,
            do_return_all_layers=True)

对Transformer内部，逐层attention
1）先搞self-attention，注意有效位置的计算attention_mask。
2）再对每个位置做前向网络，加个drop层，加个layer-norm层，再加上input_layer。
3）再对每个位置做前向网络，加个drop层，加个layer-norm层，再加上attention_output。
4）输出作下层的输入，直到N层。
重要：这里根据输入query=[batch_size * seq_length, emb_size]来梳理下计算单层self-attenion过程中的维度变化。

注意：输入词的emb_size必须跟Transformer 的输出dim=-1的size一样么，必须的，因为有残差连接，必须保持维度一致。但是，head_nums_size $\times$ size_per_head = emb_size=hidden_size不用必须成立【注意，bert代码实现上是将其设为相等的】，query_layer的最后维度只需是head_nums X 任意数。靠近输出的dense包括了drop/layer-norm操作。
self-attention的矩阵计算示例

再说明下，query–>query_layer的变换，并不是echo token独享一个转换矩阵，也不是每个位置独享一个矩阵，而是query的emb_size空间–> query_layer的hidden_size空间上的维度变换。上个图说明下，[batch_size * seq_length, input_emb_size ] w^q = [batch_size* seq_length, size_per_head*head_nums]的每个部分。
重要：attention的计算示例
1）先看下单条样本时，self-attention的计算示例SelfAttentionSingle.py
2）再看下batch样本时，self-attention的计算示例SelfAttentionBatch.py
3）最后再看下batch+heads时，self-attention的计算示例SelfAttentionBatchMultiHeads.py

attention_mask的作用在于，softmax时，对非视野内的做负向大加权，使得attention-score只计算注意在可视域范围内【非补0的地方】的数值。

  if attention_mask is not None:
    # `attention_mask` = [B, 1, F, T]
    attention_mask = tf.expand_dims(attention_mask, axis=[1])
    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and -10000.0 for masked positions.
    adder = (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0
    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    attention_scores += adder
  # Normalize the attention scores to probabilities.
  # `attention_probs` = [B, N, F, T]
  attention_probs = tf.nn.softmax(attention_scores)

注意，计算完上下文信息后，要转换成[batch, seq_length, num_heads, size_per_head]，再作其他处理。

BertModel.sequence_output 是取最后attenion层的输出。
BertModel.pooled_output 取sequence_output的第一个token“CLS”的emb，然后加个连接层。

Loss Compute

Masked Language Model的loss计算

    (masked_lm_loss,
     masked_lm_example_loss, masked_lm_log_probs) = get_masked_lm_output(
         bert_config, model.get_sequence_output(), model.get_embedding_table(),
         masked_lm_positions, masked_lm_ids, masked_lm_weights)
def get_masked_lm_output(bert_config, input_tensor, output_weights, positions,
                         label_ids, label_weights):
  ## input_tensor = model.get_sequence_output(), model的最后层输出 ## [B, seq_len, emb_size]
  ## output_weights = model.get_embedding_table(), vocab_table ## [vocab_size, emb_size]
  ## positions = msked_lm_positions 遮挡词的在句子中的位置 ## [B, seq_len] ## 前几个是位置 ## 举例 [#pos1, #pos3, #pos10, 0, 0, 0]
  ## label_ids = masked_lm_ids ## 遮挡词的ID ## [B, seq_len] ## 前几个是ID ## 举例 [119, 301, 911, 0, 0, 0]
  ## label_weights = masked_lm_weights ## 遮挡词权重=1，非遮挡词权重=0 ## [B, seq_len] ## 举例 [1, 1, 1, 0, 0, 0]
  ##               = pdding([1.0] * len(masked_lm_ids))
  """Get loss and log probs for the masked LM."""
  input_tensor = gather_indexes(input_tensor, positions)
  ## gather_indexes也会将补0取出来，所以总tensor_size是不变的, [B*seq_len, emb_size]
  ## 注意后续的计算，都是在这个尺寸上进行的 ##
  ## 又单独加了层输出dense，并增加了vocab_emb_table的乘积 + bias ##

注意：label_weight在最后计算总loss时，乘上，只计算有遮挡的位置的loss。

Next Sentence Predict的loss计算

    (next_sentence_loss, next_sentence_example_loss,
     next_sentence_log_probs) = get_next_sentence_output(
         bert_config, model.get_pooled_output(), next_sentence_labels)
def get_next_sentence_output(bert_config, input_tensor, labels):
  ## input_tensor = model.get_pooled_output() ## 模型最后输出层的第一个token-"CLS"的emb ##
  ## labels = next_sentence_labels ## 
  """Get loss and log probs for the next sentence prediction."""
  ## 剩下的就是添加个dense层，二元分类 计算loss ##

两个loss加和作总的损失

    total_loss = masked_lm_loss + next_sentence_loss

参考

1. https://github.com/google-research/bert