Function op_decoder

examples/tensorflow/decoder/utils/decoder.py:329–441

Run the decoder transformer layer by FasterTransformer. Args: inputs: A tf.Tensor with shape [batch_size * beam_width, 1, hidden_dimension]. The inputs tensor of encoder. The rank must be 3. memory_tensor: A tf.tensor with shape [batch_size * beam_width,

(inputs,
               memory_tensor,
               memory_sequence_length,
               op_self_cache,
               op_mem_cache,
               psuedo_input,
               var_dict,
               decoder_args,
               step,
               sequence_lengths)

    return self_cache, mem_cache

def op_decoder(inputs,
               memory_tensor,
               memory_sequence_length,
               op_self_cache,
               op_mem_cache,
               psuedo_input,
               var_dict,
               decoder_args,
               step,
               sequence_lengths):
    '''
    Run the decoder transformer layer by FasterTransformer.
    
    Args:
        inputs: A tf.Tensor with shape [batch_size * beam_width, 1, hidden_dimension].
                The inputs tensor of encoder. The rank must be 3.
        memory_tensor: A tf.tensor with shape [batch_size * beam_width, max(memory_sequence_length), encoder_hidden_dimension]. 
                       The results of encoder transformer layer. The rank must be 3. 
                       Note that it must be extended by beam_width times
        memory_sequence_length: A tf.Tensor with shape [batch_size * beam_width], type tf.int. 
                                The length of each sentence of results of encoder. 
                                Note that it must be extended by beam_width times
        op_self_cache: A tf.Tensor with shape [num_layer, 2, None, batch_size * beam_width, hidden_dimension]. 
                       The cache space to store the keys and values of first attention layer in each step.
        op_mem_cache: A tf.Tensor with shape [num_layer, 2, batch_size * beam_width, max(memory_sequence_length) hidden_dimension]. 
                      The cache space to store the keys and values of second attention layer.
                      Since they are same in each step, it is only need to compute them in first time. 
        psuedo_input: A tf.Tensor or null list. 
                      Put the decoder results of TensorFlow when running the TensorFlow decoder and FasterTransformer
                      decoder in one model. This prevents the race condition. 
                      It is useless when only run the FasterTransformer decoder.
        decoder_args: The arguments for decoder. The details are in the class "TransformerArgument" of common.py
        var_dict: A dict of tf.Tensor or numpy array. The variables for decoder. 
                      They can be either some tensor or some numpy array. 

    Outputs:
        outputs: A tf.Tensor with shape [batch_size * beam_width, 1, hidden_dimension].
                 The results of decoder.
    '''
    
    '''
    If fuse_qkv == True, this means that the computation of q, k, v in decoder are fused in one convolution. 
    
    Therefore, we need to split them and then passing into the decoder op. The split will bring additional overhead, 
    especially when the batch size is small because the computation time is short. 
    
    However, because most of the pretrained model on network fuse the qkv, so we fuse them as default.
    '''

    decoder_op_module = tf.load_op_library(os.path.join('./lib/libtf_decoder.so'))

    use_batch_major_op_cache, _ = get_op_cache_config(decoder_args.size_per_head, decoder_args.dtype)
    if use_batch_major_op_cache == False:
        op_self_cache = tf.contrib.framework.nest.map_structure(
                                                    lambda s: tf.concat([s, tf.zeros([decoder_args.num_layer, 1,
                                                                    tf.shape(memory_tensor)[0],
                                                                    decoder_args.hidden_dim], dtype=decoder_args.dtype)], axis=1),
                                                    op_self_cache )