def __init__(self, hparams=None):
EncoderBase.__init__(self, hparams)
with tf.variable_scope(self.variable_scope):
if self._hparams.initializer:
tf.get_variable_scope().set_initializer(
layers.get_initializer(self._hparams.initializer))
self.multihead_attention_list = []
self.poswise_networks = []
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
with tf.variable_scope('attention'):
mh_attn = MultiheadAttentionEncoder(
self._hparams.multihead_attention)
self.multihead_attention_list.append(mh_attn)
if self._hparams.dim != mh_attn.hparams.output_dim:
raise ValueError(
'The "dim" in the hparams of '
'"multihead_attention" should be equal to the '
'"dim" of TransformerEncoder')
pw_net = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
final_dim = pw_net.hparams.layers[-1]['kwargs']['units']
if self._hparams.dim != final_dim:
raise ValueError(
'The output dimenstion of '
'"poswise_feedforward" should be equal '
'to the "dim" of TransformerEncoder.')
self.poswise_networks.append(pw_net)
def __init__(self, vocab_size=None, output_layer=None, hparams=None):
ModuleBase.__init__(self, hparams)
with tf.variable_scope(self.variable_scope):
if self._hparams.initializer:
tf.get_variable_scope().set_initializer(
layers.get_initializer(self._hparams.initializer))
# Make the output layer
self._output_layer, self._vocab_size = _make_output_layer(
output_layer, vocab_size, self._hparams.output_layer_bias,
self.variable_scope)
# Make attention and poswise networks
self.multihead_attentions = {'self_att': [], 'encdec_att': []}
self.poswise_networks = []
for i in range(self._hparams.num_blocks):
layer_name = 'layer_{}'.format(i)
with tf.variable_scope(layer_name):
with tf.variable_scope("self_attention"):
multihead_attention = MultiheadAttentionEncoder(
self._hparams.multihead_attention)
self.multihead_attentions['self_att'].append(
multihead_attention)
if self._hparams.dim != \
multihead_attention.hparams.output_dim:
raise ValueError('The output dimenstion of '
'MultiheadEncoder should be equal '
'to the dim of TransformerDecoder')
with tf.variable_scope('encdec_attention'):
multihead_attention = MultiheadAttentionEncoder(
self._hparams.multihead_attention)
self.multihead_attentions['encdec_att'].append(
multihead_attention)
if self._hparams.dim != \
multihead_attention.hparams.output_dim:
raise ValueError('The output dimenstion of '
'MultiheadEncoder should be equal '
'to the dim of TransformerDecoder')
pw_net = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
final_dim = pw_net.hparams.layers[-1]['kwargs']['units']
if self._hparams.dim != final_dim:
raise ValueError(
'The output dimenstion of '
'"poswise_feedforward" should be equal '
'to the "dim" of TransformerDecoder.')
self.poswise_networks.append(pw_net)
# Built in _build()
self.context = None
self.context_sequence_length = None
self.embedding = None
self._helper = None
self._cache = None
self.max_decoding_length = None
def test_feedforward(self):
"""Tests feed-forward.
"""
hparams = {
"layers": [{
"type": "torch.nn.Linear",
"kwargs": {
"in_features": 32,
"out_features": 64
}
}, {
"type": "torch.nn.Linear",
"kwargs": {
"in_features": 64,
"out_features": 128
}
}]
}
nn = FeedForwardNetwork(hparams=hparams)
self.assertEqual(len(nn.layers), len(hparams["layers"]))
_ = nn(torch.ones(64, 16, 32))
self.assertEqual(len(nn.trainable_variables),
len(hparams["layers"]) * 2)
def __init__(self, embedding, hparams=None):
ModuleBase.__init__(self, hparams)
with tf.variable_scope(self.variable_scope):
if self._hparams.initializer:
tf.get_variable_scope().set_initializer(
layers.get_initializer(self._hparams.initializer))
if self._hparams.position_embedder_type == 'sinusoids':
self.position_embedder = SinusoidsPositionEmbedder(
self._hparams.position_embedder_hparams)
else:
self.position_embedder = PositionEmbedder(
position_size=self._hparams.position_size,
hparams=self._hparams.position_embedder_hparams)
self._embedding = embedding
self._vocab_size = self._embedding.get_shape().as_list()[0]
self.output_layer = \
self._build_output_layer(shape_list(self._embedding)[-1])
self.multihead_attentions = {'self_att': [], 'encdec_att': []}
self.poswise_networks = []
for i in range(self._hparams.num_blocks):
layer_name = 'layer_{}'.format(i)
with tf.variable_scope(layer_name):
with tf.variable_scope("self_attention"):
multihead_attention = MultiheadAttentionEncoder(
self._hparams.multihead_attention)
self.multihead_attentions['self_att'].append(
multihead_attention)
# pylint: disable=protected-access
if self._hparams.dim != \
multihead_attention._hparams.output_dim:
raise ValueError('The output dimenstion of '
'MultiheadEncoder should be equal '
'to the dim of TransformerDecoder')
with tf.variable_scope('encdec_attention'):
multihead_attention = MultiheadAttentionEncoder(
self._hparams.multihead_attention)
self.multihead_attentions['encdec_att'].append(
multihead_attention)
if self._hparams.dim != \
multihead_attention._hparams.output_dim:
raise ValueError('The output dimenstion of '
'MultiheadEncoder should be equal '
'to the dim of TransformerDecoder')
poswise_network = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
if self._hparams.dim != \
poswise_network._hparams.layers[-1]['kwargs']['units']:
raise ValueError('The output dimenstion of '
'FeedForwardNetwork should be equal '
'to the dim of TransformerDecoder')
self.poswise_networks.append(poswise_network)
def __init__(self,
vocab_size=None,
output_layer=None,
tau=None,
hparams=None):
EncoderBase.__init__(self, hparams)
with tf.variable_scope(self.variable_scope):
if self._hparams.initializer:
tf.get_variable_scope().set_initializer(
layers.get_initializer(self._hparams.initializer))
# Make the output layer
self._output_layer, self._vocab_size = _make_output_layer(
output_layer, vocab_size, self._hparams.output_layer_bias,
self.variable_scope)
# Make attention and poswise networks
self.graph_multihead_attention_list = []
self.poswise_networks = []
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
with tf.variable_scope('attention'):
mh_attn = GraphMultiheadAttentionEncoder(
self._hparams.graph_multihead_attention)
self.graph_multihead_attention_list.append(mh_attn)
if self._hparams.dim != mh_attn.hparams.output_dim:
raise ValueError(
'The "dim" in the hparams of '
'"multihead_attention" should be equal to the '
'"dim" of CrossGraphTransformerFixedLengthDecoder'
)
pw_net = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
final_dim = pw_net.hparams.layers[-1]['kwargs']['units']
if self._hparams.dim != final_dim:
raise ValueError(
'The output dimenstion of '
'"poswise_feedforward" should be equal '
'to the "dim" of CrossGraphTransformerFixedLengthDecoder.'
)
self.poswise_networks.append(pw_net)
self._helper = None
self._tau = tau
def __init__(self, hparams=None):
EncoderBase.__init__(self, hparams)
with tf.variable_scope(self.variable_scope):
if self._hparams.initializer:
tf.get_variable_scope().set_initializer(
layers.get_initializer(self._hparams.initializer))
if self._hparams.position_embedder_type == 'sinusoids':
self.position_embedder = SinusoidsPositionEmbedder(
self._hparams.position_embedder_hparams)
else:
self.position_embedder = PositionEmbedder(
position_size=self._hparams.position_size,
hparams=self._hparams.position_embedder_hparams)
# pylint: disable=protected-access
if self._hparams.dim != \
self.position_embedder._hparams.dim:
raise ValueError('"dim" in '
'TransformerEncoder hparams must be equal '
'to "dim" in its '
'position_embedder_hparams.')
self.multihead_attention_list = []
self.poswise_networks = []
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
with tf.variable_scope('attention'):
multihead_attention = MultiheadAttentionEncoder(
self._hparams.multihead_attention)
self.multihead_attention_list.append(
multihead_attention)
# pylint: disable=protected-access
if self._hparams.dim != \
multihead_attention._hparams.output_dim:
raise ValueError('The "dim" in the hparams of '
'multihead_attention should be equal '
'to the "dim" of TransformerEncoder')
poswise_network = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
# pylint: disable=protected-access
if self._hparams.dim != \
poswise_network._hparams.layers[-1]['kwargs']['units']:
# poswise_network._hparams.layers[-1]['units']:
raise ValueError('The "units" in the "kwargs" of '
'FeedForwardNetwork should be equal '
'to the "dim" of TransformerEncoder')
self.poswise_networks.append(poswise_network)
def test_feedforward(self):
"""Tests feed-forward.
"""
hparams = {
"layers": [{
"type": "Dense",
}, {
"type": "Dense",
}]
}
nn = FeedForwardNetwork(hparams=hparams)
self.assertEqual(len(nn.layers), len(hparams["layers"]))
_ = nn(tf.ones([64, 16, 16]))
self.assertEqual(len(nn.trainable_variables),
len(hparams["layers"]) * 2)
self.assertEqual(len(nn.layer_outputs), len(hparams["layers"]))
def __init__(self, hparams=None):
EncoderBase.__init__(self, hparams)
with tf.variable_scope(self.variable_scope):
if self._hparams.initializer:
tf.get_variable_scope().set_initializer(
layers.get_initializer(self._hparams.initializer))
self.position_embedder = \
SinusoidsPositionEmbedder(
self._hparams.position_embedder_hparams)
self.multihead_attention_list = []
self.poswise_networks = []
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
with tf.variable_scope('self_attention'):
multihead_attention = MultiheadAttentionEncoder(
self._hparams.multihead_attention)
self.multihead_attention_list.append(
multihead_attention)
# pylint: disable=protected-access
if self._hparams.dim != \
multihead_attention._hparams.output_dim:
raise ValueError('The output dimenstion of'
'MultiheadEncoder should be equal'
'to the dim of TransformerEncoder')
poswise_network = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
# pylint: disable=protected-access
if self._hparams.dim != \
poswise_network._hparams.layers[-1]['kwargs']['units']:
# poswise_network._hparams.layers[-1]['units']:
raise ValueError('The output dimenstion of'
'FeedForwardNetwork should be equal'
'to the dim of TransformerEncoder')
self.poswise_networks.append(poswise_network)
def _self_attention_stack(self,
inputs,
memory,
decoder_self_attention_bias=None,
memory_attention_bias=None,
cache=None,
mode=None):
"""Stacked multihead attention module.
"""
inputs = tf.layers.dropout(inputs,
rate=self._hparams.embedding_dropout,
training=is_train_mode(mode))
if cache is not None:
memory_attention_bias = \
cache['memory_attention_bias']
else:
assert decoder_self_attention_bias is not None
x = inputs
for i in range(self._hparams.num_blocks):
layer_name = 'layer_{}'.format(i)
layer_cache = cache[layer_name] if cache is not None else None
with tf.variable_scope(layer_name):
with tf.variable_scope("self_attention"):
selfatt_output = attn.multihead_attention(
queries=layers.layer_normalize(x),
memory=None,
memory_attention_bias=decoder_self_attention_bias,
num_units=self._hparams.dim,
num_heads=self._hparams.num_heads,
dropout_rate=self._hparams.attention_dropout,
cache=layer_cache,
scope="multihead_attention",
)
x = x + tf.layers.dropout(
selfatt_output,
rate=self._hparams.residual_dropout,
training=is_train_mode(mode),
)
if memory is not None:
with tf.variable_scope('encdec_attention'):
encdec_output = attn.multihead_attention(
queries=layers.layer_normalize(x),
memory=memory,
memory_attention_bias=memory_attention_bias,
num_units=self._hparams.dim,
num_heads=self._hparams.num_heads,
dropout_rate=self._hparams.attention_dropout,
scope="multihead_attention"
)
x = x + tf.layers.dropout(encdec_output, \
rate=self._hparams.residual_dropout, \
training=is_train_mode(mode))
poswise_network = FeedForwardNetwork( \
hparams=self._hparams['poswise_feedforward'])
with tf.variable_scope(poswise_network.variable_scope):
sub_output = tf.layers.dropout(
poswise_network(layers.layer_normalize(x)),
rate=self._hparams.residual_dropout,
training=is_train_mode(mode),
)
x = x + sub_output
return layers.layer_normalize(x)
def __init__(self,
vocab_size: Optional[int] = None,
output_layer: Optional[Union[nn.Module, torch.Tensor]] = None,
hparams: Optional[HParams] = None):
super().__init__(
0,
vocab_size, # dummy value for input_size
input_time_major=False,
output_time_major=False,
hparams=hparams)
self._input_size = self._hparams.dim
self._output_layer, self._vocab_size = _make_output_layer(
output_layer, vocab_size, self._input_size,
self._hparams.output_layer_bias)
self.self_attns = nn.ModuleList()
self.self_attn_layer_norm = nn.ModuleList()
self.enc_dec_attns = nn.ModuleList()
self.end_dec_attn_layer_norm = nn.ModuleList()
self.poswise_networks = nn.ModuleList()
self.poswise_layer_norm = nn.ModuleList()
if self._hparams.use_gpt_config:
eps = 1e-5
else:
eps = 1e-12
for _ in range(self._hparams.num_blocks):
attn_module = MultiheadAttentionEncoder(
self._input_size, self._hparams.multihead_attention)
if self._hparams.dim != attn_module.output_size:
raise ValueError("The output dimension of "
"MultiheadEncoder should be equal "
"to the dim of TransformerDecoder")
self.self_attns.append(attn_module)
self.self_attn_layer_norm.append(
nn.LayerNorm(self._input_size, eps=eps))
attn_module = MultiheadAttentionEncoder(
self._input_size, self._hparams.multihead_attention)
if self._hparams.dim != attn_module.output_size:
raise ValueError("The output dimension of "
"MultiheadEncoder should be equal "
"to the dim of TransformerDecoder")
self.enc_dec_attns.append(attn_module)
self.end_dec_attn_layer_norm.append(
nn.LayerNorm(self._input_size, eps=eps))
poswise_network = FeedForwardNetwork(
hparams=self._hparams.poswise_feedforward)
if (poswise_network.hparams.layers[-1]['kwargs']['out_features'] !=
self._hparams.dim):
raise ValueError("The output dimension of "
"FeedForwardNetwork should be equal "
"to the dim of TransformerDecoder")
self.poswise_networks.append(poswise_network)
self.poswise_layer_norm.append(
nn.LayerNorm(self._input_size, eps=eps))
self.final_layer_norm = nn.LayerNorm(self._input_size, eps=eps)
self.embed_dropout = nn.Dropout(self._hparams.embedding_dropout)
self.residual_dropout = nn.Dropout(self._hparams.residual_dropout)
if self._hparams.initializer:
# TODO: This might be different to what TensorFlow does
initialize = layers.get_initializer(self._hparams.initializer)
assert initialize is not None
# Do not re-initialize LayerNorm modules.
for name, param in self.named_parameters():
if name.split(
".")[-1] == "weight" and "layer_norm" not in name:
initialize(param)
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 word embeddings of input sequences. Note that
the embedding dimension `dim` must equal "dim" in
:attr:`hparams`.
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.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 = inputs * self._hparams.dim**0.5
inputs = mask_sequences(inputs, sequence_length, tensor_rank=3)
_, lengths, _ = shape_list(inputs)
inputs_padding = 1 - tf.sequence_mask(
sequence_length, tf.shape(inputs)[1], dtype=tf.float32)
ignore_padding = attn.attention_bias_ignore_padding(inputs_padding)
encoder_self_attention_bias = ignore_padding
pos_embeds = self.position_embedder(lengths, self._hparams.dim)
input_embedding = inputs + pos_embeds
x = tf.layers.dropout(input_embedding,
rate=self._hparams.embedding_dropout,
training=is_train_mode(mode))
pad_remover = utils.transformer_utils.PadRemover(inputs_padding)
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
with tf.variable_scope('self_attention'):
selfatt_output = attn.multihead_attention(
queries=layers.layer_normalize(x),
memory=None,
memory_attention_bias=encoder_self_attention_bias,
num_heads=self._hparams.num_heads,
dropout_rate=self._hparams.attention_dropout,
num_units=self._hparams.dim,
scope='multihead_attention')
x = x + tf.layers.dropout(
selfatt_output,
rate=self._hparams.residual_dropout,
training=is_train_mode(mode),
)
poswise_network = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
with tf.variable_scope(poswise_network.variable_scope):
y = layers.layer_normalize(x)
original_shape = shape_list(y)
y = tf.reshape(y, [-1, self._hparams.dim])
y = tf.expand_dims(pad_remover.remove(y), axis=0)
# [1, batch_size*seq_length, hidden_dim]
sub_output = tf.layers.dropout(
poswise_network(y),
rate=self._hparams.residual_dropout,
training=is_train_mode(mode))
sub_output = tf.reshape(pad_remover.restore(tf.squeeze(\
sub_output, axis=0)), original_shape \
)
x = x + sub_output
encoder_output = layers.layer_normalize(x)
if not self._built:
self._add_internal_trainable_variables()
self._built = True
return encoder_output
def __init__(self, hparams=None):
EncoderBase.__init__(self, hparams)
self._input_size = self._hparams.dim
self.self_attns = nn.ModuleList()
if not self._hparams.use_bert_config:
self.self_attn_layer_norm = nn.ModuleList()
self.poswise_networks = nn.ModuleList()
self.poswise_layer_norm = nn.ModuleList()
self.output_layer_norm = nn.ModuleList()
if self._hparams.use_bert_config:
# In TensorFlow, eps for LayerNorm is 1e-12 by default.
eps = 1e-12
else:
# In PyTorch, eps for LayerNorm is 1e-6 by default.
eps = 1e-6
for _ in range(self._hparams.num_blocks):
mh_attn = MultiheadAttentionEncoder(
self._input_size, self._hparams.multihead_attention)
self.self_attns.append(mh_attn)
if not self._hparams.use_bert_config:
self.self_attn_layer_norm.append(
nn.LayerNorm(self._input_size, eps=eps))
if self._hparams.dim != mh_attn.hparams.output_dim:
raise ValueError(
'The "dim" in the hparams of '
'"multihead_attention" should be equal to the '
'"dim" of TransformerEncoder')
pw_net = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
final_dim = pw_net.hparams.layers[-1]['kwargs']['out_features']
if self._hparams.dim != final_dim:
raise ValueError('The output dimenstion of '
'"poswise_feedforward" should be equal '
'to the "dim" of TransformerEncoder.')
self.poswise_networks.append(pw_net)
self.poswise_layer_norm.append(
nn.LayerNorm(self._input_size, eps=eps))
if self._hparams.use_bert_config:
self.output_layer_norm.append(
nn.LayerNorm(self._input_size, eps=eps))
self.embed_dropout = nn.Dropout(p=self._hparams.embedding_dropout)
self.residual_dropout = nn.Dropout(p=self._hparams.residual_dropout)
if self._hparams.use_bert_config:
self.input_normalizer = nn.LayerNorm(self._input_size, eps=eps)
else:
self.final_layer_normalizer = nn.LayerNorm(self._input_size,
eps=eps)
if self._hparams.initializer:
initialize = layers.get_initializer(self._hparams.initializer)
assert initialize is not None
# Do not re-initialize LayerNorm modules.
for name, param in self.named_parameters():
if name.split(
'.')[-1] == 'weight' and 'layer_norm' not in name:
initialize(param)
def _build(self, inputs, mode=None):
"""Encodes the inputs with transformer encoder.
Args:
inputs: A 2D Tensor of shape `[batch_size, max_time]`
mode(optional): A tensor taking value in
:tf_main:`tf.estimator.ModeKeys <estimator/ModeKeys>`
"""
encoder_padding = tf.to_float(tf.equal(inputs, 0))
#pylint:disable=too-many-locals
self.enc = tf.nn.embedding_lookup(self._embedding, inputs)
_, _, channels = shape_list(self.enc)
if self._hparams.multiply_embedding_mode == 'sqrt_depth':
self.enc = self.enc * channels**0.5
ignore_padding = attentions.attention_bias_ignore_padding(
encoder_padding)
encoder_self_attention_bias = ignore_padding
encoder_decoder_attention_bias = ignore_padding
if self.target_symbol_embedding:
emb_target_space = tf.reshape(self.target_symbol_embedding,
[1, 1, -1])
self.enc = self.enc + emb_target_space
lengths = shape_list(self.enc)[1]
channels = shape_list(self.enc)[2]
pos_embeds = self.position_embedder(lengths, channels)
input_embedding = self.enc + pos_embeds
x = tf.layers.dropout(input_embedding,
rate=self._hparams.embedding_dropout,
training=context.global_mode_train())
pad_remover = utils.transformer_utils.PadRemover(encoder_padding)
for i in range(self._hparams.num_blocks):
with tf.variable_scope("layer_{}".format(i)):
with tf.variable_scope('self_attention'):
selfatt_output = attentions.multihead_attention(
queries=layers.layer_normalize(x),
memory=None,
memory_attention_bias=encoder_self_attention_bias,
num_heads=self._hparams.num_heads,
dropout_rate=self._hparams.attention_dropout,
num_units=self._hparams.num_units,
scope='multihead_attention')
x = x + tf.layers.dropout(
selfatt_output,
rate=self._hparams.residual_dropout,
training=context.global_mode_train())
poswise_network = FeedForwardNetwork(
hparams=self._hparams['poswise_feedforward'])
with tf.variable_scope(poswise_network.variable_scope):
y = layers.layer_normalize(x)
original_shape = shape_list(y)
y = tf.reshape(y, [-1, self._hparams.num_units])
y = tf.expand_dims(pad_remover.remove(y), axis=0)
#[1, batch_size*seq_length, hidden_dim]
sub_output = tf.layers.dropout(
poswise_network(y),
rate=self._hparams.residual_dropout,
training=context.global_mode_train())
sub_output = tf.reshape(pad_remover.restore(tf.squeeze(\
sub_output, axis=0)), original_shape \
)
x = x + sub_output
self.stack_output = x
encoder_output = layers.layer_normalize(x)
if not self._built:
self._add_internal_trainable_variables()
self._built = True
return encoder_output, encoder_decoder_attention_bias