Function tf_bert

examples/tensorflow/bert/utils/bert.py:159–252

Run the bert transformer layer by TensorFlow. Args: inputs: A tf.Tensor with shape [batch_size, seq_len, hidden_dimension]. The inputs tensor of encoder. The rank must be 3. encoder_args: The arguments for encoder. The details are in the class

(input_tensor,
            encoder_args,
            attention_mask=None,
            intermediate_act_fn=gelu,
            initializer_range=0.02)

def tf_bert(input_tensor,
            encoder_args,
            attention_mask=None,
            intermediate_act_fn=gelu,
            initializer_range=0.02):
    '''
    Run the bert transformer layer by TensorFlow.
    
    Args:
        inputs: A tf.Tensor with shape [batch_size, seq_len, hidden_dimension]. 
                The inputs tensor of encoder. The rank must be 3. 
        encoder_args: The arguments for encoder. The details are in the class 
                      "TransformerArgument" of common.py
        attention_mask: A tf.Tensor. The attention mask for self attention.
        intermediate_act_fn: A callable function.  
                             The activation function in the FFN. It is gelu in BERT. 
        initializer_range: A float value.     
                           The range of initializer for all weights.
        
    Outputs:
        outputs: A tf.Tensor with shape [batch_size, seq_len, hidden_dimension].
                 The results of encoder.
    '''
    
    if encoder_args.hidden_dim % encoder_args.head_num != 0:
        raise ValueError(
            "The hidden size (%d) is not a multiple of the number of attention "
            "heads (%d)" % (encoder_args.hidden_dim, encoder_args.head_num))

    input_shape = get_shape_list(input_tensor, expected_rank=3)
    batch_size = input_shape[0]
    seq_length = input_shape[1]

    prev_output = reshape_to_matrix(input_tensor)

    for layer_idx in range(encoder_args.num_layer):
        with tf.variable_scope("layer_%d" % layer_idx, reuse=tf.AUTO_REUSE):
            layer_input = prev_output
            with tf.variable_scope("attention"):
                with tf.variable_scope("self"):
                    attention_head = attention_layer(
                        from_tensor=layer_input,
                        to_tensor=layer_input,
                        attention_mask=attention_mask,
                        num_attention_heads=encoder_args.head_num,
                        size_per_head=encoder_args.size_per_head,
                        initializer_range=initializer_range,
                        do_return_2d_tensor=True,
                        batch_size=batch_size,
                        from_seq_length=seq_length,
                        to_seq_length=seq_length,
                        tf_datatype=encoder_args.dtype)
                    attention_output = attention_head

                with tf.variable_scope("output"):
                    attention_output = tf.layers.dense(
                        attention_output,
                        encoder_args.hidden_dim,