def _transform(self, inputs, encoder_self_attention_bias, **kwargs): """ Encodes the inputs. Args: inputs: A Tensor, [batch_size, timesteps, d_model] encoder_self_attention_bias: A Tensor, FLOAT_MIN for padding, 0 for non-padding, [batch_size, 1, 1, timesteps]. **kwargs: Returns: A Tensor, the transformed hidden state of TransformerEncoder, [batch_size, timesteps, d_model]. """ input_padding = attention_bias_to_padding(encoder_self_attention_bias) pad_remover = PadRemover(input_padding) x = dropout_wrapper(inputs, self.params["layer_prepostprocess_dropout_keep_prob"]) encoder_self_attention_scores = [] for layer in range(self.params["num_layers"]): with tf.variable_scope("layer_%d" % layer): with tf.variable_scope("self_attention"): # self attention layer w_y, y = self._self_attention_layers[layer].build( query=None, memory=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), memory_bias=encoder_self_attention_bias) encoder_self_attention_scores.append(w_y) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("ffn"): y = transformer_ffn_layer( x=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), filter_size=self.params["num_filter_units"], output_size=self.params["num_hidden_units"], pad_remover=pad_remover, dropout_relu_keep_prob=self.params["dropout_relu_keep_prob"]) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) x = layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) return x, encoder_self_attention_scores
def _transform(self, inputs, encoder_self_attention_bias, **kwargs): """ Encodes the inputs. Args: inputs: A Tensor, [batch_size, timesteps, d_model] encoder_self_attention_bias: A Tensor, FLOAT_MIN for padding, 0 for non-padding, [batch_size, 1, 1, timesteps]. **kwargs: Returns: A Tensor, the transformed hidden state of TransformerEncoder, [batch_size, timesteps, d_model]. """ input_padding = attention_bias_to_padding(encoder_self_attention_bias) pad_remover = PadRemover(input_padding) x = dropout_wrapper(inputs, self.params["layer_prepostprocess_dropout_keep_prob"]) encoder_self_attention_scores = [] for layer in range(self.params["num_layers"]): with tf.variable_scope("layer_%d" % layer): with tf.variable_scope("self_attention"): # self attention layer w_y, y = self._self_attention_layers[layer].build( query=None, memory=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), memory_bias=encoder_self_attention_bias) encoder_self_attention_scores.append(w_y) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("ffn"): y = transformer_ffn_layer( x=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), filter_size=self.params["num_filter_units"], output_size=self.params["num_hidden_units"], pad_remover=pad_remover, dropout_relu_keep_prob=self.params["dropout_relu_keep_prob"]) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) x = layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) return x, encoder_self_attention_scores
def _transform(self, decoder_inputs, cache): """ Decodes one step Args: decoder_input: The decoder input for this timestep, A Tensor, with shape [batch_size, timesteps, dmodel]. Note that when mode==INFER, timesteps=1. cache: A dict containing decoding states at previous timestep, attention values and attention length. Returns: A transformed Tensor. """ # [batch_size, max_len_src, dim] encdec_attention_values = cache["memory"] # [batch_size, 1, 1, max_len_src] encdec_attention_bias = cache["memory_bias"] decoder_self_attention_scores = [] encdec_attention_scores = [] # decoder_self_attention_bias: [1, 1, max_len_trg, max_len_trg] decoder_self_attention_bias = attention_bias_lower_triangle( tf.shape(decoder_inputs)[1]) x = dropout_wrapper(decoder_inputs, self.params["layer_prepostprocess_dropout_keep_prob"]) for layer in range(self.params["num_layers"]): layer_name = "layer_{}".format(layer) layer_cache = None if cache["decoding_states"] is None \ else cache["decoding_states"][layer_name] selfatt_cache = None if layer_cache is None \ else layer_cache["self_attention"] encdecatt_cache = None if layer_cache is None \ else layer_cache["encdec_attention"] with tf.variable_scope("layer_%d" % layer): with tf.variable_scope("self_attention"): # self attention layer w_y, y = self._self_attention_layers[layer].build( query=None, memory=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), memory_bias=decoder_self_attention_bias, cache=selfatt_cache) # [batch_size, num_heads, length_q, length_k] decoder_self_attention_scores.append(w_y) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("encdec_attention"): # encoder-decoder attention w_y, y = self._encdec_attention_layers[layer].build( query=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), memory=encdec_attention_values, memory_bias=encdec_attention_bias, cache=encdecatt_cache) # [batch_size, num_heads, length_q, length_k] encdec_attention_scores.append(w_y) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("ffn"): y = transformer_ffn_layer( x=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), filter_size=self.params["num_filter_units"], output_size=self.params["num_hidden_units"], pad_remover=None, dropout_relu_keep_prob=self.params["dropout_relu_keep_prob"]) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) x = layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) return x, decoder_self_attention_scores, encdec_attention_scores
def _transform(self, decoder_inputs, decoding_params): """ Decodes one step Args: decoder_input: The decoder input for this timestep, an instance of `tf.Tensor`, [batch_size, timesteps, dmodel]. decoding_params: The same as `decoding_params` returned from `prepare()` function. Returns: A Tensor, the transformed hidden state of TransformerDecoder. """ # [batch_size, max_len_src, dim] encdec_attention_values = decoding_params[0] # [batch_size, ] # encdec_attention_length = decoding_params[1] # [batch_size, 1, 1, max_len_src] encdec_attention_bias = decoding_params[2] # decoder_self_attention_bias: [1, 1, max_len_trg, max_len_trg] decoder_self_attention_bias = attention_bias_lower_triangle( tf.shape(decoder_inputs)[1]) x = dropout_wrapper( decoder_inputs, self.params["layer_prepostprocess_dropout_keep_prob"]) for layer in range(self.params["num_layers"]): with tf.variable_scope("layer_%d" % layer): with tf.variable_scope("self_attention"): # self attention layer w_y, y = multihead_attention_layer( params=self.params["selfattention.params"], mode=self.mode, query_antecedent=None, memory_antecedent=layer_preprocess( x=x, process_sequence=self. params["layer_preprocess_sequence"], dropout_keep_prob=self. params["layer_prepostprocess_dropout_keep_prob"]), memory_bias=decoder_self_attention_bias) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self. params["layer_postprocess_sequence"], dropout_keep_prob=self. params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("encdec_attention"): # encoder-decoder attention w_y, y = multihead_attention_layer( params=self.params["attention.params"], mode=self.mode, query_antecedent=layer_preprocess( x=x, process_sequence=self. params["layer_preprocess_sequence"], dropout_keep_prob=self. params["layer_prepostprocess_dropout_keep_prob"]), memory_antecedent=encdec_attention_values, memory_bias=encdec_attention_bias) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self. params["layer_postprocess_sequence"], dropout_keep_prob=self. params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("ffn"): y = transformer_ffn_layer( x=layer_preprocess( x=x, process_sequence=self. params["layer_preprocess_sequence"], dropout_keep_prob=self. params["layer_prepostprocess_dropout_keep_prob"]), filter_size=self.params["num_filter_units"], output_size=self.params["num_hidden_units"], pad_remover=None, dropout_relu_keep_prob=self. params["dropout_relu_keep_prob"]) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self. params["layer_postprocess_sequence"], dropout_keep_prob=self. params["layer_prepostprocess_dropout_keep_prob"]) return layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self. params["layer_prepostprocess_dropout_keep_prob"])
def _transform(self, decoder_inputs, cache, pad_remover=None): """ Decodes one step Args: decoder_inputs: The decoder input for this timestep, A Tensor, with shape [batch_size, timesteps, dmodel]. Note that when mode==INFER, timesteps=1. cache: A dict containing decoding states at previous timestep, attention values and attention length. pad_remover: An expert_utils.PadRemover object tracking the padding positions. If provided, the padding is removed before applying the convolution, and restored afterward. This can give a significant speedup (says Google's tensor2tensor code). Returns: A transformed Tensor. """ # [batch_size, max_len_src, dim] encdec_attention_values = cache["memory"] # [batch_size, 1, 1, max_len_src] encdec_attention_bias = cache["memory_bias"] decoder_self_attention_scores = [] encdec_attention_scores = [] # decoder_self_attention_bias: [1, 1, max_len_trg, max_len_trg] decoder_self_attention_bias = attention_bias_lower_triangle( tf.shape(decoder_inputs)[1]) x = dropout_wrapper(decoder_inputs, self.params["layer_prepostprocess_dropout_keep_prob"]) for layer in range(self.params["num_layers"]): layer_name = "layer_{}".format(layer) layer_cache = None if cache["decoding_states"] is None \ else cache["decoding_states"][layer_name] selfatt_cache = None if layer_cache is None \ else layer_cache["self_attention"] encdecatt_cache = None if layer_cache is None \ else layer_cache["encdec_attention"] with tf.variable_scope("layer_%d" % layer): with tf.variable_scope("self_attention"): # self attention layer w_y, y = self._self_attention_layers[layer].build( query=None, memory=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), memory_bias=decoder_self_attention_bias, cache=selfatt_cache) # [batch_size, num_heads, length_q, length_k] decoder_self_attention_scores.append(w_y) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("encdec_attention"): # encoder-decoder attention w_y, y = self._encdec_attention_layers[layer].build( query=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), memory=encdec_attention_values, memory_bias=encdec_attention_bias, cache=encdecatt_cache) # [batch_size, num_heads, length_q, length_k] encdec_attention_scores.append(w_y) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) with tf.variable_scope("ffn"): y = transformer_ffn_layer( x=layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]), filter_size=self.params["num_filter_units"], output_size=self.params["num_hidden_units"], pad_remover=pad_remover, dropout_relu_keep_prob=self.params["dropout_relu_keep_prob"]) # apply dropout, layer norm, residual x = layer_postprocessing( x=y, previous_x=x, process_sequence=self.params["layer_postprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) x = layer_preprocess( x=x, process_sequence=self.params["layer_preprocess_sequence"], dropout_keep_prob=self.params["layer_prepostprocess_dropout_keep_prob"]) return x, decoder_self_attention_scores, encdec_attention_scores