【Bert】（六）句子关系判断--源码解析（bert基础模型）

论文：https://arxiv.org/pdf/1810.04805.pdf

官方代码：GitHub - google-research/bert: TensorFlow code and pre-trained models for BERT

这里主要将bert模型分为了两个部分：bert基础模型+bert后处理模型。bert基础模型可以理解为一个通用模型，用来提取数据的特征。bert后处理模型指的是在获取到特征后，将特征运用来解决特定任务。

bert基础模型在不同任务上是一样的，bert后处理模型在不同任务上会有不同的结构。

(total_loss, per_example_loss, logits, probabilities) = create_model(
        bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
        num_labels, use_one_hot_embeddings)

上述代码是由bert基础模型+bert后处理模型所构成。

下面具体讲解bert基础模型，入口代码为：

  model = modeling.BertModel(
      config=bert_config,
      is_training=is_training,
      input_ids=input_ids,
      input_mask=input_mask,
      token_type_ids=segment_ids,
      use_one_hot_embeddings=use_one_hot_embeddings)

1、配置类（BertConfig）

这部分代码主要定义了BERT模型的一些默认参数，另外包括了一些文件处理函数。

class BertConfig(object):
  """BERT模型的配置类."""

  def __init__(self,
               vocab_size,
               hidden_size=768,
               num_hidden_layers=12,
               num_attention_heads=12,
               intermediate_size=3072,
               hidden_act="gelu",
               hidden_dropout_prob=0.1,
               attention_probs_dropout_prob=0.1,
               max_position_embeddings=512,
               type_vocab_size=16,
               initializer_range=0.02):

    self.vocab_size = vocab_size
    self.hidden_size = hidden_size
    self.num_hidden_layers = num_hidden_layers
    self.num_attention_heads = num_attention_heads
    self.hidden_act = hidden_act
    self.intermediate_size = intermediate_size
    self.hidden_dropout_prob = hidden_dropout_prob
    self.attention_probs_dropout_prob = attention_probs_dropout_prob
    self.max_position_embeddings = max_position_embeddings
    self.type_vocab_size = type_vocab_size
    self.initializer_range = initializer_range

  @classmethod
  def from_dict(cls, json_object):
    """Constructs a `BertConfig` from a Python dictionary of parameters."""
    config = BertConfig(vocab_size=None)
    for (key, value) in six.iteritems(json_object):
      config.__dict__[key] = value
    return config

  @classmethod
  def from_json_file(cls, json_file):
    """Constructs a `BertConfig` from a json file of parameters."""
    with tf.gfile.GFile(json_file, "r") as reader:
      text = reader.read()
    return cls.from_dict(json.loads(text))

  def to_dict(self):
    """Serializes this instance to a Python dictionary."""
    output = copy.deepcopy(self.__dict__)
    return output

  def to_json_string(self):
    """Serializes this instance to a JSON string."""
    return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

参数具体含义

vocab_size：词表大小（不需要主动设置，是看词表中有多少词）
hidden_size：隐藏层神经元数（在uncased_L-12_H-768_A-12/bert_config.json中）
num_hidden_layers：Transformer encoder中的隐藏层数（在uncased_L-12_H-768_A-12/bert_config.json中）
num_attention_heads：multi-head attention 的head数（在uncased_L-12_H-768_A-12/bert_config.json中）
intermediate_size：encoder的“中间”隐层神经元数（例如feed-forward layer）
hidden_act：隐藏层激活函数
hidden_dropout_prob：隐层dropout率
attention_probs_dropout_prob：注意力部分的dropout
max_position_embeddings：最大位置编码
type_vocab_size：token_type_ids的词典大小
initializer_range：truncated_normal_initializer初始化方法的stdev

这里要注意一点，可能刚看的时候对type_vocab_size这个参数会有点不理解，其实就是在next sentence prediction任务里的Segment A和 Segment B。在下载的bert_config.json文件里也有说明，默认值应该为2。

参考这个Issue​github.com/google-research/bert/issues/16

2、获取词向量（Embedding_lookup）

对于输入word_ids，返回embedding table。可以选用one-hot或者tf.gather()

def embedding_lookup(input_ids,                        # word_id：【batch_size, seq_length】
                     vocab_size,
                     embedding_size=128,
                     initializer_range=0.02,
                     word_embedding_name="word_embeddings",
                     use_one_hot_embeddings=False):

  # 该函数默认输入的形状为【batch_size, seq_length, input_num】
  # 如果输入为2D的【batch_size, seq_length】，则扩展到【batch_size, seq_length, 1】
  if input_ids.shape.ndims == 2:
    input_ids = tf.expand_dims(input_ids, axis=[-1])

  embedding_table = tf.get_variable(
      name=word_embedding_name,
      shape=[vocab_size, embedding_size],
      initializer=create_initializer(initializer_range))

  flat_input_ids = tf.reshape(input_ids, [-1])    #【batch_size*seq_length*input_num】
  if use_one_hot_embeddings:
    one_hot_input_ids = tf.one_hot(flat_input_ids, depth=vocab_size)
    output = tf.matmul(one_hot_input_ids, embedding_table)
  else:    # 按索引取值
    output = tf.gather(embedding_table, flat_input_ids)

  input_shape = get_shape_list(input_ids)

  # output：[batch_size, seq_length, num_inputs]
  # 转成:[batch_size, seq_length, num_inputs*embedding_size]
  output = tf.reshape(output,
                        input_shape[0:-1] + [input_shape[-1] * embedding_size])
  return (output, embedding_table)

参数具体含义

input_ids：word id 【batch_size, seq_length】
vocab_size：embedding词表大小
embedding_size：embedding维度
initializer_range：embedding初始化范围（这个值表示正太分布的标准差，均值默认为0）
word_embedding_name：embeddding table命名
use_one_hot_embeddings：是否使用one-hotembedding
Return：【batch_size, seq_length, embedding_size】

这里原作者使用的是gather。

这里的操作是，生成一个[vocab_size=30522, embedding_size=768]的正态分布向量，也就是说词汇表30522的每个词都被一个768维的向量所表示，这个gather，其实就是根据索引值将对应的768维的向量挑选出来表示该单词。所以最终返回的维度为【batch_size, seq_length, embedding_size】。

这里说一下我对将词转化为向量的理解：

将输入的一段话向量化我想到的有三种方式：（1）直接使用词汇所对应的索引值；（2）根据词汇的索引值使用one-hot方式；（3）每一个词汇都用固定维度的随机向量表示。

方式（1）使用索引值，假设词汇表的索引值取值范围[0,9999]，那么如此大的范围直接送入神经网络中，容易被激活函数抑制，而且不容易拟合。那将区间范围压缩至0-1，那这临近单词之间的距离只有0.0001，实在太小了，也不利于学习。所以这种方式被否决。

方式（2）将词根据索引值进行one-hot表示，具体one-hot表示啥请自行百度。这里我们发现似乎他解决了方式（1）所带来的问题，但是两个词之间的余弦距离为0，也就是他不能表示词与词之间的关联性。而且当词汇表的数量增加时，所有词的one-hot表示都需要更新，这种方式总体不是很友好。

方式（3）将词用一个固定维度的向量表示，这个向量可以进行学习更新，那么这个向量之间的余弦距离不是0，也就是说模型可能会学到词与词之间的关联性信息，而且出现新增词汇时，已经存在词的向量可以不用更新，很是友好。

3、词向量的后续处理（embedding_postprocessor）

我们知道BERT模型的输入有三部分：token embedding ，segment embedding以及position embedding。上一节中我们只获得了token embedding，这部分代码对其完善信息，正则化，dropout之后输出最终embedding。
注意，在Transformer论文中的position embedding是由sin/cos函数生成的固定的值，而在这里代码实现中是跟普通word embedding一样随机生成的，可以训练的。作者这里这样选择的原因可能是BERT训练的数据比Transformer那篇大很多，完全可以让模型自己去学习。

def embedding_postprocessor(input_tensor,                # [batch_size, seq_length, embedding_size]
                            use_token_type=False,
                            token_type_ids=None,
                            token_type_vocab_size=16,        # 一般是2
                            token_type_embedding_name="token_type_embeddings",
                            use_position_embeddings=True,
                            position_embedding_name="position_embeddings",
                            initializer_range=0.02,
                            max_position_embeddings=512,    #最大位置编码，必须大于等于max_seq_len
                            dropout_prob=0.1):

  input_shape = get_shape_list(input_tensor, expected_rank=3)   #【batch_size,seq_length,embedding_size】
  batch_size = input_shape[0]
  seq_length = input_shape[1]
  width = input_shape[2]

  output = input_tensor

  # Segment position信息
  if use_token_type:
    if token_type_ids is None:
      raise ValueError("`token_type_ids` must be specified if"
                       "`use_token_type` is True.")
    token_type_table = tf.get_variable(
        name=token_type_embedding_name,
        shape=[token_type_vocab_size, width],
        initializer=create_initializer(initializer_range))
    # 由于token-type-table比较小，所以这里采用one-hot的embedding方式加速
    flat_token_type_ids = tf.reshape(token_type_ids, [-1])
    one_hot_ids = tf.one_hot(flat_token_type_ids, depth=token_type_vocab_size)
    token_type_embeddings = tf.matmul(one_hot_ids, token_type_table)
    token_type_embeddings = tf.reshape(token_type_embeddings,
                                       [batch_size, seq_length, width])
    output += token_type_embeddings

  # Position embedding信息
  if use_position_embeddings:
    # 确保seq_length小于等于max_position_embeddings
    assert_op = tf.assert_less_equal(seq_length, max_position_embeddings)
    with tf.control_dependencies([assert_op]):
      full_position_embeddings = tf.get_variable(
          name=position_embedding_name,
          shape=[max_position_embeddings, width],
          initializer=create_initializer(initializer_range))

      # 这里position embedding是可学习的参数，[max_position_embeddings, width]
      # 但是通常实际输入序列没有达到max_position_embeddings
      # 所以为了提高训练速度，使用tf.slice取出句子长度的embedding
      position_embeddings = tf.slice(full_position_embeddings, [0, 0],
                                     [seq_length, -1])
      num_dims = len(output.shape.as_list())

      # word embedding之后的tensor是[batch_size, seq_length, width]
      # 因为位置编码是与输入内容无关，它的shape总是[seq_length, width]
      # 我们无法把位置Embedding加到word embedding上
      # 因此我们需要扩展位置编码为[1, seq_length, width]
      # 然后就能通过broadcasting加上去了。
      position_broadcast_shape = []
      for _ in range(num_dims - 2):
        position_broadcast_shape.append(1)
      position_broadcast_shape.extend([seq_length, width])
      position_embeddings = tf.reshape(position_embeddings,
                                       position_broadcast_shape)
      output += position_embeddings

  output = layer_norm_and_dropout(output, dropout_prob)
  return output

如上代码中segment position部分，token_type_ids表示一段话中的第一句还是第二句，第一句用0表示，第二句用1表示，其他用0表示。将token_type_ids使用onehot方式变为[batchsize*128, 2]，然后构建一个可学习变量[2, 768]，两者进行矩阵乘法，得到[batchsize*128, 768]，然后reshape成[batchsize, 128, 768]。但是具体为什么使用这种操作，我不知道原理。

position_embeddings是随机了一个[512,768]的张量，由于句子设定的长度为128，所以选取前128为张量，得到[128,768]的丈量，为了与token_position能进行相加，将其扩展为[1,128,768]，之后使用广播broadcast的方式进行相加即可。

这里的2个embedding信息为什么这么操作似乎原理都不是很能明白，难道模型知道哪个是哪个吗，真能学到我们赋予他的信息吗？

4、构造attention_mask

from_tensor：[batchsize, 128]，表示一段话中每个单词对应词汇表中单词的索引值。

to_mask：[batchsize, 128]，表示一段话中哪些是真正的句子，真正的句子为1，padding的句子为0。

mask：[batchsize, 128，128]。

构造的attention_mask可以看出来没有变量不需要学习，全都是常量的计算，那说明只经过一次运算该值就被确定，那么表明该方法的返回值是一种辅助张量，旨在提供额外信息。

其实这里的作用就是为了方便attention_layer中的一个softmax计算，紧紧是为了让padding出来的语句在经过softmax后变成一个非常小的值，不影响后续的模型的学习。

def create_attention_mask_from_input_mask(from_tensor, to_mask):
  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)

  broadcast_ones = tf.ones(
      shape=[batch_size, from_seq_length, 1], dtype=tf.float32)

  mask = broadcast_ones * to_mask

  return mask

5、注意力层（attention layer）

这部分代码是multi-head attention的实现，主要来自《Attention is all you need》这篇论文。考虑key-query-value形式的attention，输入的from_tensor当做是query， to_tensor当做是key和value，当两者相同的时候即为self-attention。关于attention更详细的介绍可以转到理解Attention机制原理及模型。

def attention_layer(from_tensor,   # 【batch_size, from_seq_length, from_width】
                    to_tensor,        #【batch_size, to_seq_length, to_width】
                    attention_mask=None,        #【batch_size,from_seq_length, to_seq_length】
                    num_attention_heads=1,        # attention head numbers
                    size_per_head=512,            # 每个head的大小
                    query_act=None,                # query变换的激活函数
                    key_act=None,                # key变换的激活函数
                    value_act=None,                # value变换的激活函数
                    attention_probs_dropout_prob=0.0,        # attention层的dropout
                    initializer_range=0.02,                    # 初始化取值范围
                    do_return_2d_tensor=False,                # 是否返回2d张量。
#如果True，输出形状【batch_size*from_seq_length,num_attention_heads*size_per_head】
#如果False，输出形状【batch_size, from_seq_length, num_attention_heads*size_per_head】
                    batch_size=None,                        #如果输入是3D的，
#那么batch就是第一维，但是可能3D的压缩成了2D的，所以需要告诉函数batch_size 
                    from_seq_length=None,                    # 同上
                    to_seq_length=None):                    # 同上

  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])    #[batch_size,  num_attention_heads, seq_length, width]
    return output_tensor

  from_shape = get_shape_list(from_tensor, expected_rank=[2, 3])
  to_shape = get_shape_list(to_tensor, expected_rank=[2, 3])

  if len(from_shape) != len(to_shape):
    raise ValueError(
        "The rank of `from_tensor` must match the rank of `to_tensor`.")

  if len(from_shape) == 3:
    batch_size = from_shape[0]
    from_seq_length = from_shape[1]
    to_seq_length = to_shape[1]
  elif len(from_shape) == 2:
    if (batch_size is None or from_seq_length is None or to_seq_length is None):
      raise ValueError(
          "When passing in rank 2 tensors to attention_layer, the values "
          "for `batch_size`, `from_seq_length`, and `to_seq_length` "
          "must all be specified.")

  # 为了方便备注shape，采用以下简写:
  #   B = batch size (number of sequences)
  #   F = `from_tensor` sequence length
  #   T = `to_tensor` sequence length
  #   N = `num_attention_heads`
  #   H = `size_per_head`

  # 把from_tensor和to_tensor压缩成2D张量
  from_tensor_2d = reshape_to_matrix(from_tensor)        # 【B*F, hidden_size】
  to_tensor_2d = reshape_to_matrix(to_tensor)            # 【B*T, hidden_size】

  # 将from_tensor输入全连接层得到query_layer
  # `query_layer` = [B*F, N*H]
  query_layer = tf.layers.dense(
      from_tensor_2d,
      num_attention_heads * size_per_head,
      activation=query_act,
      name="query",
      kernel_initializer=create_initializer(initializer_range))

  # 将from_tensor输入全连接层得到query_layer
  # `key_layer` = [B*T, N*H]
  key_layer = tf.layers.dense(
      to_tensor_2d,
      num_attention_heads * size_per_head,
      activation=key_act,
      name="key",
      kernel_initializer=create_initializer(initializer_range))

  # 同上
  # `value_layer` = [B*T, N*H]
  value_layer = tf.layers.dense(
      to_tensor_2d,
      num_attention_heads * size_per_head,
      activation=value_act,
      name="value",
      kernel_initializer=create_initializer(initializer_range))

  # query_layer转成多头：[B*F, N*H]==>[B, F, N, H]==>[B, N, F, H]
  query_layer = transpose_for_scores(query_layer, batch_size,
                                     num_attention_heads, from_seq_length,
                                     size_per_head)

  # key_layer转成多头：[B*T, N*H] ==> [B, T, N, H] ==> [B, N, T, H]
  key_layer = transpose_for_scores(key_layer, batch_size, num_attention_heads,
                                   to_seq_length, size_per_head)

  # 将query与key做点积，然后做一个scale，公式可以参见原始论文
  # `attention_scores` = [B, N, F, T]
  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` = [B, 1, F, T]
    attention_mask = tf.expand_dims(attention_mask, axis=[1])

    # 如果attention_mask里的元素为1，则通过下面运算有（1-1）*-10000，adder就是0
    # 如果attention_mask里的元素为0，则通过下面运算有（1-0）*-10000，adder就是-10000
    adder = (1.0 - tf.cast(attention_mask, tf.float32)) * -10000.0

    # 我们最终得到的attention_score一般不会很大，
    #所以上述操作对mask为0的地方得到的score可以认为是负无穷
    attention_scores += adder

  # 负无穷经过softmax之后为0，就相当于mask为0的位置不计算attention_score
  # `attention_probs` = [B, N, F, T]
  attention_probs = tf.nn.softmax(attention_scores)

  # 对attention_probs进行dropout，这虽然有点奇怪，但是Transforme原始论文就是这么做的
  attention_probs = dropout(attention_probs, attention_probs_dropout_prob)

  # `value_layer` = [B, T, N, H]
  value_layer = tf.reshape(
      value_layer,
      [batch_size, to_seq_length, num_attention_heads, size_per_head])

  # `value_layer` = [B, N, T, H]
  value_layer = tf.transpose(value_layer, [0, 2, 1, 3])

  # `context_layer` = [B, N, F, H]
  context_layer = tf.matmul(attention_probs, value_layer)

  # `context_layer` = [B, F, N, H]
  context_layer = tf.transpose(context_layer, [0, 2, 1, 3])

  if do_return_2d_tensor:
    # `context_layer` = [B*F, N*H]
    context_layer = tf.reshape(
        context_layer,
        [batch_size * from_seq_length, num_attention_heads * size_per_head])
  else:
    # `context_layer` = [B, F, N*H]
    context_layer = tf.reshape(
        context_layer,
        [batch_size, from_seq_length, num_attention_heads * size_per_head])

  return context_layer

总结一下，attention layer的主要流程：

• 对输入的tensor进行形状校验，提取batch_size、from_seq_length 、to_seq_length
• 输入如果是3d张量则转化成2d矩阵
• from_tensor作为query， to_tensor作为key和value，经过一层全连接层后得到query_layer、key_layer 、value_layer
• 将上述张量通过transpose_for_scores转化成multi-head（就是所谓的多头注意力机制，仅仅是为了提升模型的变化程度，也变相增加了模型的能力，这里为啥能增加模型能力我不是很懂。）
• 根据论文公式计算attention_score以及attention_probs（注意attention_mask的trick）： 
• 将得到的attention_probs与value相乘，返回2D或3D张量，其形状为

  [batch_size * from_seq_length, num_attention_heads * size_per_head]

这段代码主要就是实现如下一个运算过程，其他的操作都是在变换形状方便计算，或者是抑制无效信息，使得学习更为准确和简单。

6、Transformer

接下来的代码就是大名鼎鼎的Transformer的核心代码了，可以认为是"Attention is All You Need"原始代码重现。可以参见原始论文和原始代码。

def transformer_model(input_tensor,                        # 【batch_size, seq_length, hidden_size】
                      attention_mask=None,                # 【batch_size, seq_length, seq_length】
                      hidden_size=768,
                      num_hidden_layers=12,
                      num_attention_heads=12,
                      intermediate_size=3072,
                      intermediate_act_fn=gelu,            # feed-forward层的激活函数
                      hidden_dropout_prob=0.1,
                      attention_probs_dropout_prob=0.1,
                      initializer_range=0.02,
                      do_return_all_layers=False):

  # 这里注意，因为最终要输出hidden_size， 我们有num_attention_head个区域，
  # 每个head区域有size_per_head多的隐层
  # 所以有 hidden_size = num_attention_head * size_per_head
  if hidden_size % num_attention_heads != 0:
    raise ValueError(
        "The hidden size (%d) is not a multiple of the number of attention "
        "heads (%d)" % (hidden_size, num_attention_heads))

  attention_head_size = int(hidden_size / num_attention_heads)
  input_shape = get_shape_list(input_tensor, expected_rank=3)
  batch_size = input_shape[0]
  seq_length = input_shape[1]
  input_width = input_shape[2]

  # 因为encoder中有残差操作，所以需要shape相同
  if input_width != hidden_size:
    raise ValueError("The width of the input tensor (%d) != hidden size (%d)" %
                     (input_width, hidden_size))

  # reshape操作在CPU/GPU上很快，但是在TPU上很不友好
  # 所以为了避免2D和3D之间的频繁reshape，我们把所有的3D张量用2D矩阵表示
  prev_output = reshape_to_matrix(input_tensor)

  all_layer_outputs = []
  for layer_idx in range(num_hidden_layers):
    with tf.variable_scope("layer_%d" % layer_idx):
      layer_input = prev_output

      with tf.variable_scope("attention"):
      # multi-head attention
        attention_heads = []
        with tf.variable_scope("self"):
        # self-attention
          attention_head = attention_layer(
              from_tensor=layer_input,
              to_tensor=layer_input,
              attention_mask=attention_mask,
              num_attention_heads=num_attention_heads,
              size_per_head=attention_head_size,
              attention_probs_dropout_prob=attention_probs_dropout_prob,
              initializer_range=initializer_range,
              do_return_2d_tensor=True,
              batch_size=batch_size,
              from_seq_length=seq_length,
              to_seq_length=seq_length)
          attention_heads.append(attention_head)

        attention_output = None
        if len(attention_heads) == 1:
          attention_output = attention_heads[0]
        else:
          # 如果有多个head，将他们拼接起来
          attention_output = tf.concat(attention_heads, axis=-1)

        # 对attention的输出进行线性映射, 目的是将shape变成与input一致
        # 然后dropout+residual+norm
        with tf.variable_scope("output"):
          attention_output = tf.layers.dense(
              attention_output,
              hidden_size,
              kernel_initializer=create_initializer(initializer_range))
          attention_output = dropout(attention_output, hidden_dropout_prob)
          attention_output = layer_norm(attention_output + layer_input)

      # feed-forward
      with tf.variable_scope("intermediate"):
        intermediate_output = tf.layers.dense(
            attention_output,
            intermediate_size,
            activation=intermediate_act_fn,
            kernel_initializer=create_initializer(initializer_range))

      # 对feed-forward层的输出使用线性变换变回‘hidden_size’
      # 然后dropout + residual + norm
      with tf.variable_scope("output"):
        layer_output = tf.layers.dense(
            intermediate_output,
            hidden_size,
            kernel_initializer=create_initializer(initializer_range))
        layer_output = dropout(layer_output, hidden_dropout_prob)
        layer_output = layer_norm(layer_output + attention_output)
        prev_output = layer_output
        all_layer_outputs.append(layer_output)

  if do_return_all_layers:
    final_outputs = []
    for layer_output in all_layer_outputs:
      final_output = reshape_from_matrix(layer_output, input_shape)
      final_outputs.append(final_output)
    return final_outputs
  else:
    final_output = reshape_from_matrix(prev_output, input_shape)
    return final_output

配上下图一同使用效果更佳，因为BERT里只有encoder，所有decoder没有姓名

7、函数入口（init）

BertModel类的构造函数，有了上面几节的铺垫，我们就可以来实现BERT模型了。

def __init__(self,
               config,                            # BertConfig对象
               is_training,                        
               input_ids,                        # 【batch_size, seq_length】
               input_mask=None,                    # 【batch_size, seq_length】
               token_type_ids=None,                # 【batch_size, seq_length】
               use_one_hot_embeddings=False,    # 是否使用one-hot；否则tf.gather()
               scope=None):

    config = copy.deepcopy(config)
    if not is_training:
      config.hidden_dropout_prob = 0.0
      config.attention_probs_dropout_prob = 0.0

    input_shape = get_shape_list(input_ids, expected_rank=2)
    batch_size = input_shape[0]
    seq_length = input_shape[1]
    # 不做mask，即所有元素为1
    if input_mask is None:
      input_mask = tf.ones(shape=[batch_size, seq_length], dtype=tf.int32)

    if token_type_ids is None:
      token_type_ids = tf.zeros(shape=[batch_size, seq_length], dtype=tf.int32)

    with tf.variable_scope(scope, default_name="bert"):
      with tf.variable_scope("embeddings"):
        # word embedding 
        (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)

        # 添加position embedding和segment embedding
        # layer norm + dropout
        self.embedding_output = embedding_postprocessor(
            input_tensor=self.embedding_output,
            use_token_type=True,
            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,
            position_embedding_name="position_embeddings",
            initializer_range=config.initializer_range,
            max_position_embeddings=config.max_position_embeddings,
            dropout_prob=config.hidden_dropout_prob)

      with tf.variable_scope("encoder"):

        # input_ids是经过padding的word_ids： [25, 120, 34, 0, 0]
        # input_mask是有效词标记：            [1, 1, 1, 0, 0]
        attention_mask = create_attention_mask_from_input_mask(
            input_ids, input_mask)

        # 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)

      # `self.sequence_output`是最后一层的输出，shape为【batch_size, seq_length, hidden_size】
      self.sequence_output = self.all_encoder_layers[-1]

      # ‘pooler’部分将encoder输出【batch_size, seq_length, hidden_size】
      # 转成【batch_size, hidden_size】
      with tf.variable_scope("pooler"):
        # 取最后一层的第一个时刻[CLS]对应的tensor， 对于分类任务很重要
        # sequence_output[:, 0:1, :]得到的是[batch_size, 1, hidden_size]
        # 我们需要用squeeze把第二维去掉
        first_token_tensor = tf.squeeze(self.sequence_output[:, 0:1, :], axis=1)
        # 然后再加一个全连接层，输出仍然是[batch_size, hidden_size]
        self.pooled_output = tf.layers.dense(
            first_token_tensor,
            config.hidden_size,
            activation=tf.tanh,
            kernel_initializer=create_initializer(config.initializer_range))

self.sequence_output是最后一层的输出，shape为【batch_size, seq_length, hidden_size】 

总结一哈

有了以上对源码的深入了解之后，我们在使用BertModel的时候就会更加得心应手。举个模型使用的简单栗子：

# 假设输入已经经过分词变成word_ids. shape=[2, 3]
input_ids = tf.constant([[31, 51, 99], [15, 5, 0]])
input_mask = tf.constant([[1, 1, 1], [1, 1, 0]])
# segment_emebdding. 表示第一个样本前两个词属于句子1，后一个词属于句子2.
# 第二个样本的第一个词属于句子1， 第二次词属于句子2，第三个元素0表示padding
token_type_ids = tf.constant([[0, 0, 1], [0, 2, 0]])

# 创建BertConfig实例
config = modeling.BertConfig(vocab_size=32000, hidden_size=512,
         num_hidden_layers=8, num_attention_heads=6, intermediate_size=1024)

# 创建BertModel实例
model = modeling.BertModel(config=config, is_training=True,
     input_ids=input_ids, input_mask=input_mask, token_type_ids=token_type_ids)


label_embeddings = tf.get_variable(...)
#得到最后一层的第一个Token也就是[CLS]向量表示，可以看成是一个句子的embedding
pooled_output = model.get_pooled_output()
logits = tf.matmul(pooled_output, label_embeddings)

在BERT模型构建这一块的主要流程：

• 对输入序列进行Embedding（三个），接下去就是‘Attention is all you need’的内容了
• 简单一点就是将embedding输入transformer得到输出结果
• 详细一点就是embedding --> N *【multi-head attention --> Add(Residual) &Norm--> Feed-Forward --> Add(Residual) &Norm】
• 哈，是不是很简单~
• 源码中还有一些其他的辅助函数，不是很难理解，这里就不再啰嗦。