天天看點

Bert核心代碼解讀

前面已經介紹了如何先測試一個bert任務,對這方面還不了解的可以看一下前面的部落格。

BERT 最主要的模型實作部分---BertModel,代碼位于

  • modeling.py 子產品

為了便于了解,下面的代碼中的batch_size假設成8,seq_length長度是128,每個詞編碼後的向量緯度是768。

配置類(BertConfig)

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:隐藏層神經元數
  • num_hidden_layers:Transformer encoder 中的隐藏層數
  • *num_attention_heads:*multi-head attention 的 head 數
  • 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。

函數入口(init)

上面看完了類的定義後,這個檔案的主要代碼在BertModel裡面,我們來看一下BertModel 類的構造函數。

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,                        # 8x128
            vocab_size=config.vocab_size,               # 模型中的詞表
            embedding_size=config.hidden_size,          # 想要把詞映射成多少緯度,tf官網給出的緯度是768。
            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))
           

接下來我們将會按照init代碼從上到下解讀主要函數的代碼。

擷取詞向量(Embedding_lookup)

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】  比如8X128
  # 如果輸入為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(    # 詞映射矩陣,30522,768,在詞表内進行查找
      name=word_embedding_name,
      shape=[vocab_size, embedding_size],     # 30522,768
      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) # 一個batch裡所有的映射結果

  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)    # (8,128,768)
           

參數定義:

  • input_ids:word id 【batch_size, seq_length】
  • vocab_size:embedding 詞表
  • embedding_size:embedding 次元
  • initializer_range:embedding 初始化範圍
  • word_embedding_name:embeddding table 命名
  • use_one_hot_embeddings:是否使用 one-hotembedding
  • Return:【batch_size, seq_length, embedding_size】

該子產品是将一個詞轉換成向量的子產品,它的輸入是batch_size*seq_length,輸出是batch_size, seq_length, embedding_size。這個embedding_size是每個詞所映射成向量後的緯度。在編碼向量的時候此處就是在bert預訓練好的模型中查找詞的向量。bert_model.ckpt.data-00000-of-00001裡面存放的是訓練好的詞向量,vocab.txt裡面存放的是詞表。具體存放如下所示:

Bert核心代碼解讀

位置編碼(embedding_postprocessor)

我們知道 BERT 模型的輸入有三部分:

token embedding

 ,

segment embedding

以及

position embedding

。在 Transformer 論文中的

position embedding

是由 sin/cos 函數生成的固定的值,而在這裡代碼實作中是跟普通 word embedding 一樣随機生成的,可以訓練的。

Bert核心代碼解讀
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( # (2,768),2的意思是隻有兩種結果,第一句和第二句,第一句用0表示,第二句用1表示
        name=token_type_embedding_name,
        shape=[token_type_vocab_size, width],
        initializer=create_initializer(initializer_range))
    # This vocab will be small so we always do one-hot here, since it is always
    # faster for a small vocabulary.
    # 由于token-type-table比較小,是以這裡采用one-hot的embedding方式加速
    flat_token_type_ids = tf.reshape(token_type_ids, [-1])  # 對8x128=1024個詞都要找到segment position資訊,每個詞有兩種可能性
    one_hot_ids = tf.one_hot(flat_token_type_ids, depth=token_type_vocab_size)# 1024,2和2,768做乘法
    token_type_embeddings = tf.matmul(one_hot_ids, token_type_table)
    token_type_embeddings = tf.reshape(token_type_embeddings,
                                       [batch_size, seq_length, width]) # 8,128,768
    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])    # 位置編碼緯度過大,此處為了加速隻取出有用的部分,128,768
      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,    # [1,128,768]表示位置編碼跟輸入資料無關
                                       position_broadcast_shape)
      output += position_embeddings

  output = layer_norm_and_dropout(output, dropout_prob)
  return output
           

先加上了一個segment position的資訊,然後才加Position embedding的資訊。先初始化一個(2,768)的矩陣,2的意思是隻有兩種結果,第一句和第二句,第一句用0表示,第二句用1表示。然後在做segment的時候對8x128=1024個詞都要找到segment position資訊,每個詞有兩種可能性。找的方法是用矩陣的乘法。(1024,2)和(2,768)大小的矩陣做乘法。

在加上位置編碼的時候先初始化一個位置矩陣,剛開始初始化矩陣的緯度可能比較大,假設是512,也就是說有512個位置。矩陣的大小是512*768,目的在于和詞向量長度一樣。然後取和seq_length大小一樣的部分,此處seq_length的大小是128。得到的位置編碼矩陣大小是128x768,768是保證和詞向量的緯度相同,後面将位置編碼進行擴充并且和詞向量進行相加。傳回的output是整個embedding的結果。

構造 attention_mask

該子產品大概了解作用即可,此處不再對代碼進行過多的解讀。它的輸入是一個二維的向量,輸出是一個3D的矩陣。新增的緯度作用在于讓一句話中的每一個詞編碼的向量能夠看到自己可以進行計算的向量。下面我用一個圖來解釋,假設下圖中的向量是8個句子做完embedding後的向量,後面的0代表句子的長度已結束。此時第一個句子的第一個編碼在後面做self-Attention所需要和該句子中的其他向量計算,那麼該和哪些向量計算呐?此處就是用新增加的緯度來表示需要計算的詞向量,圖中下部分是轉換成3D後新增加的一個向量來表示和哪些詞進行計算,1代表能計算,0代表不進行計算。該部分的核心代碼在create_attention_mask_from_input_mask子產品。

Bert核心代碼解讀

注意力層(attention layer)

這部分代碼是「multi-head attention」的實作,主要來自《Attention is all you need》這篇論文。考慮

key-query-value

形式的 attention,輸入的

from_tensor

當做是 query, 

to_tensor

當做是 key 和 value,當兩者相同的時候即為 self-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)  8
  #   F = `from_tensor` sequence length 128
  #   T = `to_tensor` sequence length   128
  #   N = `num_attention_heads` 12
  #   H = `size_per_head`   每個頭有64個特征
  # 把from_tensor和to_tensor壓縮成2D張量

  # 把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):
    Bert核心代碼解讀
  • 将得到的 attention_probs 與 value 相乘,傳回 2D 或 3D 張量
Bert核心代碼解讀

Transformer

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核心代碼解讀

如果對transformer理論還不是很了解的可以轉到https://blog.csdn.net/one_super_dreamer/article/details/105181690

繼續閱讀