def _layer_norm(x, scope):
return layers.layer_normalize(x, reuse=tf.AUTO_REUSE, scope=scope)
def _build(self, inputs, sequence_length, mode=None):
"""Encodes the inputs.
Args:
inputs: A 3D Tensor of shape `[batch_size, max_time, dim]`,
containing the embedding of input sequences. Note that
the embedding dimension `dim` must equal "dim" in
:attr:`hparams`. The input embedding is typically an aggregation
of word embedding and position embedding.
sequence_length: A 1D Tensor of shape `[batch_size]`. Input tokens
beyond respective sequence lengths are masked out
automatically.
mode (optional): A tensor taking value in
:tf_main:`tf.estimator.ModeKeys <estimator/ModeKeys>`,
including `TRAIN`, `EVAL`, and `PREDICT`. Used to toggle
dropout.
If `None` (default), :func:`texar.tf.global_mode` is used.
Returns:
A Tensor of shape `[batch_size, max_time, dim]` containing the
encoded vectors.
"""
# Multiply input embedding with the sqrt of its dimension for
# normalization
inputs_padding = 1 - tf.sequence_mask(
sequence_length, tf.shape(inputs)[1], dtype=tf.float32)
if self._hparams.use_bert_config:
ignore_padding = attn.attention_bias_ignore_padding(
inputs_padding, bias_value=-1e4)
else:
ignore_padding = attn.attention_bias_ignore_padding(
inputs_padding)
encoder_self_attention_bias = ignore_padding
input_embedding = inputs
if self._hparams.use_bert_config:
x = layers.layer_normalize(input_embedding)
x = tf.layers.dropout(x,
rate=self._hparams.embedding_dropout,
training=is_train_mode(mode))
else:
x = tf.layers.dropout(input_embedding,
rate=self._hparams.embedding_dropout,
training=is_train_mode(mode))
# Just to keep consistent with BERT, actually makes no difference
if self._hparams.use_bert_config:
pad_remover = None
else:
pad_remover = transformer_utils.PadRemover(inputs_padding)
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
multihead_attention = self.multihead_attention_list[i]
# trivial difference between BERT and original Transformer
if self._hparams.use_bert_config:
_queries_input = x
else:
_queries_input = layers.layer_normalize(x)
attention_output = multihead_attention(
queries=_queries_input,
memory=_queries_input,
memory_attention_bias=encoder_self_attention_bias,
mode=mode,
)
attention_output = tf.layers.dropout(
attention_output,
rate=self._hparams.residual_dropout,
training=is_train_mode(mode),
)
x = x + attention_output
with tf.variable_scope('output'):
if self._hparams.use_bert_config:
x = layers.layer_normalize(x)
y = x
else:
y = layers.layer_normalize(x)
poswise_network = self.poswise_networks[i]
with tf.variable_scope(poswise_network.variable_scope):
original_shape = shape_list(y)
y = tf.reshape(y, [-1, self._hparams.dim])
if pad_remover:
y = tf.expand_dims(pad_remover.remove(y), axis=0)
# [1, batch_size*seq_length, hidden_dim]
layer_output = poswise_network(y, mode=mode)
sub_output = tf.layers.dropout(
layer_output,
rate=self._hparams.residual_dropout,
training=is_train_mode(mode)
)
if pad_remover:
sub_output = tf.reshape(
pad_remover.restore(tf.squeeze(sub_output, axis=0)),
original_shape)
else:
sub_output = tf.reshape(sub_output, original_shape)
x = x + sub_output
if self._hparams.use_bert_config:
x = layers.layer_normalize(x)
if not self._hparams.use_bert_config:
x = layers.layer_normalize(x)
if not self._built:
self._add_internal_trainable_variables()
self._built = True
return x
def _build(self,
inputs,
memory,
sequence_length,
memory_sequence_length,
encoder_output,
adjs=None,
mode=None):
"""Encodes the inputs.
Args:
inputs: A 3D Tensor of shape `[batch_size, max_time, dim]`,
containing the embedding of input sequences. Note that
the embedding dimension `dim` must equal "dim" in
:attr:`hparams`. The input embedding is typically an aggregation
of word embedding and position embedding.
memory: A 3D Tensor of shape `[batch_size, memory_max_time, dim]`,
containing the embedding of memory sequences. Note that
the embedding dimension `dim` must equal "dim" in
:attr:`hparams`. The input embedding is typically an aggregation
of word embedding and position embedding.
sequence_length: A 1D Tensor of shape `[batch_size]`. Input tokens
beyond respective sequence lengths are masked out
automatically.
sequence_length: A 1D Tensor of shape `[batch_size]`. Memory tokens
beyond respective sequence lengths are masked out
automatically.
adjs: A 3D Tensor of shape `[batch_size, max_time, max_time]`,
containing the adjacency matrices of input sequences
encoder_output: bool. True: return encoder-like embeddings. False: return CrossGraphTransformerDecoderOutput.
mode (optional): A tensor taking value in
:tf_main:`tf.estimator.ModeKeys <estimator/ModeKeys>`,
including `TRAIN`, `EVAL`, and `PREDICT`. Used to toggle
dropout.
If `None` (default), :func:`texar.tf.global_mode` is used.
Returns:
A Tensor of shape `[batch_size, max_time, dim]` containing the
encoded vectors.
"""
# Get adjacency masks from adjs
if self._hparams.use_adj:
adj_masks = 1 - tf.cast(tf.equal(adjs, 0), dtype=tf.float32)
else:
adj_masks = None
# Multiply input embedding with the sqrt of its dimension for
# normalization
inputs_padding = 1 - tf.sequence_mask(
sequence_length, tf.shape(inputs)[1], dtype=tf.float32)
if self._hparams.use_bert_config:
ignore_padding = attn.attention_bias_ignore_padding(
inputs_padding, bias_value=-1e4)
else:
ignore_padding = attn.attention_bias_ignore_padding(inputs_padding)
encoder_self_attention_bias = ignore_padding
input_embedding = inputs # shape (batch_size, max_time, dim)
if self._hparams.use_bert_config:
x = layers.layer_normalize(input_embedding)
x = tf.layers.dropout(x,
rate=self._hparams.embedding_dropout,
training=is_train_mode(mode))
else:
x = tf.layers.dropout(input_embedding,
rate=self._hparams.embedding_dropout,
training=is_train_mode(mode))
# Just to keep consistent with BERT, actually makes no difference
if self._hparams.use_bert_config:
pad_remover = None
else:
pad_remover = utils.transformer_utils.PadRemover(inputs_padding)
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
graph_multihead_attention = self.graph_multihead_attention_list[
i]
# trivial difference between BERT and original Transformer
if self._hparams.use_bert_config:
_queries_input = x
else:
_queries_input = layers.layer_normalize(x)
attention_output = graph_multihead_attention(
queries=_queries_input,
memory=memory,
adj_masks=adj_masks,
memory_attention_bias=encoder_self_attention_bias,
mode=mode,
use_adj=self._hparams.use_adj)
attention_output = tf.layers.dropout(
attention_output,
rate=self._hparams.residual_dropout,
training=is_train_mode(mode),
)
# attention_output: weighted sum of V of memory with weights determined by querying keys of memory
x = x + attention_output
with tf.variable_scope('output'):
if self._hparams.use_bert_config:
x = layers.layer_normalize(x)
y = x
else:
y = layers.layer_normalize(x)
poswise_network = self.poswise_networks[i]
with tf.variable_scope(poswise_network.variable_scope):
original_shape = shape_list(y)
y = tf.reshape(y, [-1, self._hparams.dim])
if pad_remover:
y = tf.expand_dims(pad_remover.remove(y), axis=0)
# [1, batch_size*seq_length, hidden_dim]
layer_output = poswise_network(y, mode=mode)
sub_output = tf.layers.dropout(
layer_output,
rate=self._hparams.residual_dropout,
training=is_train_mode(mode))
if pad_remover:
sub_output = tf.reshape(pad_remover.restore(tf.squeeze(\
sub_output, axis=0)), original_shape \
)
else:
sub_output = tf.reshape(sub_output, original_shape)
x = x + sub_output
if self._hparams.use_bert_config:
x = layers.layer_normalize(x)
if not self._hparams.use_bert_config:
x = layers.layer_normalize(x)
if not self._built:
self._add_internal_trainable_variables()
self._built = True
if encoder_output:
return x
logits = self._output_layer(x)
sample_ids = tf.to_int32(tf.argmax(logits, axis=-1))
probs = ''
# probs = GumbelSoftmax(self._tau, logits=logits).sample()
# probs = tf.nn.softmax(logits / self._tau) # vanilla softmax
rets = CrossGraphTransformerFixedLengthDecoderOutput(
logits=logits, sample_id=sample_ids, probs=probs)
return rets
