def testSRU(self):
x = np.random.rand(5, 7, 3, 11)
with self.test_session() as session:
y = common_layers.sru(tf.constant(x, dtype=tf.float32))
session.run(tf.global_variables_initializer())
res = session.run(y)
self.assertEqual(res.shape, (5, 7, 3, 11))
def testSRU(self):
x = np.random.rand(5, 7, 3, 11)
with self.test_session() as session:
y = common_layers.sru(tf.constant(x, dtype=tf.float32))
session.run(tf.global_variables_initializer())
res = session.run(y)
self.assertEqual(res.shape, (5, 7, 3, 11))
def body(self, features):
hparams = self.hparams
is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
# Run the basic autoencoder part first.
basic_result, losses = super(AutoencoderAutoregressive, self).body(features)
if hparams.autoregressive_mode == "none":
assert not hparams.autoregressive_forget_base
return basic_result, losses
shape = common_layers.shape_list(basic_result)
basic1d = tf.reshape(basic_result, [shape[0], -1, shape[3]])
# During autoregressive inference, don't resample.
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
if hasattr(hparams, "sampled_basic1d_tensor"):
basic1d = hparams.sampled_basic1d_tensor
else:
hparams.sampled_basic1d_tensor = basic1d
# Prepare inputs for autoregressive modes.
if common_layers.shape_list(features["targets"])[1] == 1:
# This happens on the first step of predicitions.
assert hparams.mode == tf.estimator.ModeKeys.PREDICT
features["targets"] = tf.zeros_like(basic_result)
targets_dropout = common_layers.mix(
features["targets"], tf.zeros_like(basic_result),
hparams.bottleneck_warmup_steps, is_training,
max_prob=1.0 - hparams.autoregressive_dropout, broadcast_last=True)
# Sometimes it's useful to look at non-autoregressive evals.
if (hparams.mode == tf.estimator.ModeKeys.EVAL and
hparams.autoregressive_eval_pure_autoencoder):
targets_dropout = tf.zeros_like(basic_result)
# Now combine the basic reconstruction with shifted targets.
targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[3]])
targets_shifted = common_layers.shift_right_3d(targets1d)
concat1d = tf.concat([basic1d, targets_shifted], axis=-1)
# The forget_base hparam sets purely-autoregressive mode, no autoencoder.
if hparams.autoregressive_forget_base:
concat1d = tf.reshape(features["targets"], [shape[0], -1, shape[3]])
concat1d = common_layers.shift_right_3d(concat1d)
# The autoregressive part depends on the mode.
if hparams.autoregressive_mode == "conv3":
res = common_layers.conv1d(concat1d, shape[3], 3, padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv3")
return tf.reshape(res, shape), losses
if hparams.autoregressive_mode == "conv5":
res = common_layers.conv1d(concat1d, shape[3], 5, padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv5")
return tf.reshape(res, shape), losses
if hparams.autoregressive_mode == "sru":
res = common_layers.conv1d(concat1d, shape[3], 3, padding="LEFT",
activation=common_layers.belu,
name="autoregressive_sru_conv3")
res = common_layers.sru(res)
return tf.reshape(res, shape), losses
raise ValueError("Unsupported autoregressive mode: %s"
% hparams.autoregressive_mode)
def testSRU(self):
if tf.executing_eagerly():
return # don't run test in Eager mode
x = np.random.rand(5, 7, 3, 11)
with self.session() as session:
y = common_layers.sru(tf.constant(x, dtype=tf.float32))
session.run(tf.global_variables_initializer())
res = session.run(y)
self.assertEqual(res.shape, (5, 7, 3, 11))
def body(self, features):
hparams = self._hparams
shape = common_layers.shape_list(features["targets"])
# Run the basic autoencoder part first.
basic_result, losses = super(AutoencoderAutoregressive,
self).body(features)
# Prepare inputs for autoregressive modes.
targets_keep_prob = 1.0 - hparams.autoregressive_dropout
targets_dropout = common_layers.dropout_with_broadcast_dims(
features["targets"], targets_keep_prob, broadcast_dims=[-1])
targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[3]])
targets_shifted = common_layers.shift_right_3d(targets1d)
basic1d = tf.reshape(basic_result, [shape[0], -1, shape[3]])
concat1d = tf.concat([basic1d, targets_shifted], axis=-1)
# The forget_base hparam sets purely-autoregressive mode, no autoencoder.
if hparams.autoregressive_forget_base:
concat1d = tf.reshape(features["targets"],
[shape[0], -1, shape[3]])
concat1d = common_layers.shift_right_3d(concat1d)
# The autoregressive part depends on the mode.
if hparams.autoregressive_mode == "none":
assert not hparams.autoregressive_forget_base
return basic_result, losses
if hparams.autoregressive_mode == "conv3":
res = common_layers.conv1d(concat1d,
shape[3],
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv3")
return tf.reshape(res, shape), losses
if hparams.autoregressive_mode == "conv5":
res = common_layers.conv1d(concat1d,
shape[3],
5,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv5")
return tf.reshape(res, shape), losses
if hparams.autoregressive_mode == "sru":
res = common_layers.conv1d(concat1d,
shape[3],
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_sru_conv3")
res = common_layers.sru(res)
return tf.reshape(res, shape), losses
raise ValueError("Unsupported autoregressive mode: %s" %
hparams.autoregressive_mode)
def body(self, features):
hparams = self.hparams
is_training = hparams.mode == tf.estimator.ModeKeys.TRAIN
# Run the basic autoencoder part first.
basic_result, losses = super(AutoencoderAutoregressive, self).body(features)
if hparams.autoregressive_mode == "none":
assert not hparams.autoregressive_forget_base
return basic_result, losses
shape = common_layers.shape_list(basic_result)
basic1d = tf.reshape(basic_result, [shape[0], -1, shape[3]])
# During autoregressive inference, don't resample.
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
if hasattr(hparams, "sampled_basic1d_tensor"):
basic1d = hparams.sampled_basic1d_tensor
else:
hparams.sampled_basic1d_tensor = basic1d
# Prepare inputs for autoregressive modes.
if common_layers.shape_list(features["targets"])[1] == 1:
# This happens on the first step of predicitions.
assert hparams.mode == tf.estimator.ModeKeys.PREDICT
features["targets"] = tf.zeros_like(basic_result)
targets_dropout = common_layers.mix(
features["targets"],
tf.zeros_like(basic_result),
hparams.bottleneck_warmup_steps,
is_training,
max_prob=1.0 - hparams.autoregressive_dropout,
broadcast_last=True)
# Sometimes it's useful to look at non-autoregressive evals.
if (hparams.mode == tf.estimator.ModeKeys.EVAL and
hparams.autoregressive_eval_pure_autoencoder):
targets_dropout = tf.zeros_like(basic_result)
# Now combine the basic reconstruction with shifted targets.
targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[3]])
targets_shifted = common_layers.shift_right_3d(targets1d)
concat1d = tf.concat([basic1d, targets_shifted], axis=-1)
# The forget_base hparam sets purely-autoregressive mode, no autoencoder.
if hparams.autoregressive_forget_base:
concat1d = tf.reshape(features["targets"], [shape[0], -1, shape[3]])
concat1d = common_layers.shift_right_3d(concat1d)
# The autoregressive part depends on the mode.
if hparams.autoregressive_mode == "conv3":
res = common_layers.conv1d(
concat1d,
shape[3],
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv3")
return tf.reshape(res, shape), losses
if hparams.autoregressive_mode == "conv5":
res = common_layers.conv1d(
concat1d,
shape[3],
5,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv5")
return tf.reshape(res, shape), losses
if hparams.autoregressive_mode == "sru":
res = common_layers.conv1d(
concat1d,
shape[3],
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_sru_conv3")
res = common_layers.sru(res)
return tf.reshape(res, shape), losses
raise ValueError(
"Unsupported autoregressive mode: %s" % hparams.autoregressive_mode)
def body(self, features):
hparams = self.hparams
# Run the basic autoencoder part first.
basic_result, losses = super(AutoencoderAutoregressive, self).body(features)
if hparams.autoregressive_mode == "none":
assert not hparams.autoregressive_forget_base
return basic_result, losses
if "training" in losses:
plain_training_loss = losses.pop("training")
losses["plain"] = plain_training_loss
res_shape = common_layers.shape_list(basic_result)
vocab_size = self._problem_hparams.vocab_size["targets"]
if hasattr(self._hparams, "vocab_divisor"):
vocab_size += (-vocab_size) % self._hparams.vocab_divisor
targets = tf.one_hot(features["targets_raw"], vocab_size)
# Prepare inputs for autoregressive modes.
if common_layers.shape_list(features["targets"])[1] == 1:
# This happens on the first step of predicitions.
assert hparams.mode == tf.estimator.ModeKeys.PREDICT
targets = tf.zeros_like(basic_result)
targets = self.embed(targets)
if hparams.autoregressive_gumbel_sample:
basic_hot = self.gumbel_sample(basic_result)
else:
basic_hot = basic_result
basic_result = self.embed(basic_hot)
shape = common_layers.shape_list(basic_result)
basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
targets = tf.reshape(targets, common_layers.shape_list(basic_result))
# During autoregressive inference, don't resample.
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
if hasattr(hparams, "sampled_basic1d_tensor"):
basic1d = hparams.sampled_basic1d_tensor
else:
hparams.sampled_basic1d_tensor = basic1d
# Sometimes it's useful to look at non-autoregressive evals.
targets_dropout = targets
if (hparams.mode == tf.estimator.ModeKeys.EVAL and
hparams.autoregressive_eval_pure_autoencoder):
targets_dropout = tf.zeros_like(basic_result)
# Now combine the basic reconstruction with shifted targets.
targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[-1]])
targets_shifted = common_layers.shift_right_3d(targets1d)
concat1d = tf.concat([basic1d, targets_shifted], axis=-1)
# The forget_base hparam sets purely-autoregressive mode, no autoencoder.
if hparams.autoregressive_forget_base:
concat1d = tf.reshape(targets, [shape[0], -1, shape[-1]])
concat1d = common_layers.shift_right_3d(concat1d)
# The autoregressive part depends on the mode.
if hparams.autoregressive_mode == "conv3":
res = common_layers.conv1d(
concat1d,
hparams.hidden_size,
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv3")
res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
return tf.reshape(res, res_shape), losses
if hparams.autoregressive_mode == "conv5":
res = common_layers.conv1d(
concat1d,
hparams.hidden_size,
5,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv5")
res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
return tf.reshape(res, res_shape), losses
if hparams.autoregressive_mode == "sru":
res = common_layers.conv1d(
concat1d,
hparams.hidden_size,
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_sru_conv3")
res = common_layers.sru(res)
res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
return tf.reshape(res, res_shape), losses
raise ValueError(
"Unsupported autoregressive mode: %s" % hparams.autoregressive_mode)
def transformer_ffn_layer(x,
hparams,
pad_remover=None,
conv_padding="LEFT",
nonpadding_mask=None,
losses=None,
cache=None,
decode_loop_step=None,
readout_filter_size=0):
"""Feed-forward layer in the transformer.
Args:
x: a Tensor of shape [batch_size, length, hparams.hidden_size]
hparams: hyperparameters for model
pad_remover: an expert_utils.PadRemover object tracking the padding
positions. If provided, when using convolutional settings, the padding
is removed before applying the convolution, and restored afterward. This
can give a significant speedup.
conv_padding: a string - either "LEFT" or "SAME".
nonpadding_mask: an optional Tensor with shape [batch_size, length].
needed for convolutional layers with "SAME" padding.
Contains 1.0 in positions corresponding to nonpadding.
losses: optional list onto which to append extra training losses
cache: dict, containing tensors which are the results of previous
attentions, used for fast decoding.
decode_loop_step: An integer, step number of the decoding loop.
Only used for inference on TPU.
readout_filter_size: if it's greater than 0, then it will be used instead of
filter_size
Returns:
a Tensor of shape [batch_size, length, hparams.hidden_size]
Raises:
ValueError: If losses arg is None, but layer generates extra losses.
"""
ffn_layer = hparams.ffn_layer
relu_dropout_broadcast_dims = (
common_layers.comma_separated_string_to_integer_list(
getattr(hparams, "relu_dropout_broadcast_dims", "")))
if ffn_layer == "conv_hidden_relu":
# Backwards compatibility
ffn_layer = "dense_relu_dense"
if ffn_layer == "dense_relu_dense":
# In simple convolution mode, use `pad_remover` to speed up processing.
mlperf_log.transformer_print(
key=mlperf_log.MODEL_HP_FFN_FILTER_DENSE,
value={
"filter_size": hparams.filter_size,
"use_bias": "True",
"activation": mlperf_log.RELU
})
mlperf_log.transformer_print(
key=mlperf_log.MODEL_HP_FFN_OUTPUT_DENSE,
value={
"hidden_size": hparams.hidden_size,
"use_bias": "True",
})
mlperf_log.transformer_print(
key=mlperf_log.MODEL_HP_RELU_DROPOUT, value=hparams.relu_dropout)
if pad_remover:
original_shape = common_layers.shape_list(x)
# Collapse `x` across examples, and remove padding positions.
x = tf.reshape(x, tf.concat([[-1], original_shape[2:]], axis=0))
x = tf.expand_dims(pad_remover.remove(x), axis=0)
conv_output = common_layers.dense_relu_dense(
x,
hparams.filter_size,
hparams.hidden_size,
dropout=hparams.relu_dropout,
dropout_broadcast_dims=relu_dropout_broadcast_dims)
if pad_remover:
# Restore `conv_output` to the original shape of `x`, including padding.
conv_output = tf.reshape(
pad_remover.restore(tf.squeeze(conv_output, axis=0)), original_shape)
return conv_output
elif ffn_layer == "conv_relu_conv":
return common_layers.conv_relu_conv(
x,
readout_filter_size or hparams.filter_size,
hparams.hidden_size,
first_kernel_size=hparams.conv_first_kernel,
second_kernel_size=1,
padding=conv_padding,
nonpadding_mask=nonpadding_mask,
dropout=hparams.relu_dropout,
cache=cache,
decode_loop_step=decode_loop_step)
elif ffn_layer == "parameter_attention":
return common_attention.parameter_attention(
x, hparams.parameter_attention_key_channels or hparams.hidden_size,
hparams.parameter_attention_value_channels or hparams.hidden_size,
hparams.hidden_size, readout_filter_size or hparams.filter_size,
hparams.num_heads,
hparams.attention_dropout)
elif ffn_layer == "conv_hidden_relu_with_sepconv":
return common_layers.conv_hidden_relu(
x,
readout_filter_size or hparams.filter_size,
hparams.hidden_size,
kernel_size=(3, 1),
second_kernel_size=(31, 1),
padding="LEFT",
dropout=hparams.relu_dropout)
elif ffn_layer == "sru":
return common_layers.sru(x)
elif ffn_layer == "local_moe_tpu":
overhead = (
hparams.moe_overhead_train
if hparams.mode == tf.estimator.ModeKeys.TRAIN else
hparams.moe_overhead_eval)
ret, loss = expert_utils.local_moe_tpu(
x,
hparams.filter_size // 2,
hparams.hidden_size,
hparams.moe_num_experts,
overhead=overhead,
loss_coef=hparams.moe_loss_coef)
elif ffn_layer == "local_moe":
overhead = (
hparams.moe_overhead_train
if hparams.mode == tf.estimator.ModeKeys.TRAIN else
hparams.moe_overhead_eval)
ret, loss = expert_utils.local_moe(
x,
True,
expert_utils.ffn_expert_fn(hparams.hidden_size, [hparams.filter_size],
hparams.hidden_size),
hparams.moe_num_experts,
k=hparams.moe_k,
hparams=hparams)
losses.append(loss)
return ret
else:
assert ffn_layer == "none"
return x
def transformer_ffn_layer(x,
hparams,
customized_ffn=None,
pad_remover=None,
conv_padding="LEFT",
nonpadding_mask=None,
losses=None,
cache=None):
"""Feed-forward layer in the transformer.
Args:
x: a Tensor of shape [batch_size, length, hparams.hidden_size]
hparams: hyperparameters for model
customized_ffn: customized the ffn_layer string
pad_remover: an expert_utils.PadRemover object tracking the padding
positions. If provided, when using convolutional settings, the padding
is removed before applying the convolution, and restored afterward. This
can give a significant speedup.
conv_padding: a string - either "LEFT" or "SAME".
nonpadding_mask: an optional Tensor with shape [batch_size, length].
needed for convolutional layers with "SAME" padding.
Contains 1.0 in positions corresponding to nonpadding.
losses: optional list onto which to append extra training losses
cache: dict, containing tensors which are the results of previous
attentions, used for fast decoding.
Returns:
a Tensor of shape [batch_size, length, hparams.hidden_size]
Raises:
ValueError: If losses arg is None, but layer generates extra losses.
"""
ffn_layer = customized_ffn or hparams.ffn_layer
relu_dropout_broadcast_dims = (
common_layers.comma_separated_string_to_integer_list(
getattr(hparams, "relu_dropout_broadcast_dims", "")))
if ffn_layer == "conv_hidden_relu":
# Backwards compatibility
ffn_layer = "dense_relu_dense"
if ffn_layer == "dense_relu_dense":
# In simple convolution mode, use `pad_remover` to speed up processing.
if pad_remover:
original_shape = common_layers.shape_list(x)
# Collapse `x` across examples, and remove padding positions.
x = tf.reshape(x, tf.concat([[-1], original_shape[2:]], axis=0))
x = tf.expand_dims(pad_remover.remove(x), axis=0)
conv_output = common_layers.dense_relu_dense(
x,
hparams.filter_size,
hparams.hidden_size,
dropout=hparams.relu_dropout,
dropout_broadcast_dims=relu_dropout_broadcast_dims)
if pad_remover:
# Restore `conv_output` to the original shape of `x`, including padding.
conv_output = tf.reshape(
pad_remover.restore(tf.squeeze(conv_output, axis=0)), original_shape)
return conv_output
elif ffn_layer == "conv_relu_conv":
return common_layers.conv_relu_conv(
x,
hparams.filter_size,
hparams.hidden_size,
first_kernel_size=hparams.conv_first_kernel,
second_kernel_size=1,
padding=conv_padding,
nonpadding_mask=nonpadding_mask,
dropout=hparams.relu_dropout,
cache=cache)
elif ffn_layer == "parameter_attention":
return common_attention.parameter_attention(
x, hparams.parameter_attention_key_channels or hparams.hidden_size,
hparams.parameter_attention_value_channels or hparams.hidden_size,
hparams.hidden_size, hparams.filter_size, hparams.num_heads,
hparams.attention_dropout)
elif ffn_layer == "conv_hidden_relu_with_sepconv":
return common_layers.conv_hidden_relu(
x,
hparams.filter_size,
hparams.hidden_size,
kernel_size=(3, 1),
second_kernel_size=(31, 1),
padding="LEFT",
dropout=hparams.relu_dropout)
elif ffn_layer == "sru":
return common_layers.sru(x)
elif ffn_layer == "local_moe_tpu":
overhead = (hparams.moe_overhead_train
if hparams.mode == tf.estimator.ModeKeys.TRAIN
else hparams.moe_overhead_eval)
ret, loss = expert_utils.local_moe_tpu(
x, hparams.filter_size // 2,
hparams.hidden_size,
hparams.moe_num_experts, overhead=overhead,
loss_coef=hparams.moe_loss_coef)
if losses is None:
raise ValueError(
"transformer_ffn_layer with type local_moe_tpu must pass in "
"a losses list")
losses.append(loss)
return ret
else:
assert ffn_layer == "none"
return x
def body(self, features):
hparams = self.hparams
# Run the basic autoencoder part first.
basic_result, losses = super(AutoencoderAutoregressive, self).body(features)
if hparams.autoregressive_mode == "none":
assert not hparams.autoregressive_forget_base
return basic_result, losses
if "training" in losses:
plain_training_loss = losses.pop("training")
losses["plain"] = plain_training_loss
res_shape = common_layers.shape_list(basic_result)
vocab_size = self._problem_hparams.modality["targets"].top_dimensionality
targets = tf.one_hot(features["targets_raw"], vocab_size)
# Prepare inputs for autoregressive modes.
if common_layers.shape_list(features["targets"])[1] == 1:
# This happens on the first step of predicitions.
assert hparams.mode == tf.estimator.ModeKeys.PREDICT
targets = tf.zeros_like(basic_result)
targets = self.embed(targets)
if hparams.autoregressive_gumbel_sample:
basic_hot = self.gumbel_sample(basic_result)
else:
basic_hot = basic_result
basic_result = self.embed(basic_hot)
shape = common_layers.shape_list(basic_result)
basic1d = tf.reshape(basic_result, [shape[0], -1, shape[-1]])
targets = tf.reshape(targets, common_layers.shape_list(basic_result))
# During autoregressive inference, don't resample.
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
if hasattr(hparams, "sampled_basic1d_tensor"):
basic1d = hparams.sampled_basic1d_tensor
else:
hparams.sampled_basic1d_tensor = basic1d
# Sometimes it's useful to look at non-autoregressive evals.
targets_dropout = targets
if (hparams.mode == tf.estimator.ModeKeys.EVAL and
hparams.autoregressive_eval_pure_autoencoder):
targets_dropout = tf.zeros_like(basic_result)
# Now combine the basic reconstruction with shifted targets.
targets1d = tf.reshape(targets_dropout, [shape[0], -1, shape[-1]])
targets_shifted = common_layers.shift_right_3d(targets1d)
concat1d = tf.concat([basic1d, targets_shifted], axis=-1)
# The forget_base hparam sets purely-autoregressive mode, no autoencoder.
if hparams.autoregressive_forget_base:
concat1d = tf.reshape(targets, [shape[0], -1, shape[-1]])
concat1d = common_layers.shift_right_3d(concat1d)
# The autoregressive part depends on the mode.
if hparams.autoregressive_mode == "conv3":
res = common_layers.conv1d(
concat1d,
hparams.hidden_size,
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv3")
res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
return tf.reshape(res, res_shape), losses
if hparams.autoregressive_mode == "conv5":
res = common_layers.conv1d(
concat1d,
hparams.hidden_size,
5,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_conv5")
res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
return tf.reshape(res, res_shape), losses
if hparams.autoregressive_mode == "sru":
res = common_layers.conv1d(
concat1d,
hparams.hidden_size,
3,
padding="LEFT",
activation=common_layers.belu,
name="autoregressive_sru_conv3")
res = common_layers.sru(res)
res = tf.layers.dense(res, vocab_size, name="autoregressive_final")
return tf.reshape(res, res_shape), losses
raise ValueError(
"Unsupported autoregressive mode: %s" % hparams.autoregressive_mode)
def transformer_ffn_layer(x,
hparams,
pad_remover=None,
conv_padding="LEFT",
nonpadding_mask=None,
losses=None,
cache=None,
decode_loop_step=None,
readout_filter_size=0):
"""Feed-forward layer in the transformer.
Args:
x: a Tensor of shape [batch_size, length, hparams.hidden_size]
hparams: hyperparameters for model
pad_remover: an expert_utils.PadRemover object tracking the padding
positions. If provided, when using convolutional settings, the padding
is removed before applying the convolution, and restored afterward. This
can give a significant speedup.
conv_padding: a string - either "LEFT" or "SAME".
nonpadding_mask: an optional Tensor with shape [batch_size, length].
needed for convolutional layers with "SAME" padding.
Contains 1.0 in positions corresponding to nonpadding.
losses: optional list onto which to append extra training losses
cache: dict, containing tensors which are the results of previous
attentions, used for fast decoding.
decode_loop_step: An integer, step number of the decoding loop.
Only used for inference on TPU.
readout_filter_size: if it's greater than 0, then it will be used instead of
filter_size
Returns:
a Tensor of shape [batch_size, length, hparams.hidden_size]
Raises:
ValueError: If losses arg is None, but layer generates extra losses.
"""
ffn_layer = hparams.ffn_layer
relu_dropout_broadcast_dims = (
common_layers.comma_separated_string_to_integer_list(
getattr(hparams, "relu_dropout_broadcast_dims", "")))
if ffn_layer == "conv_hidden_relu":
# Backwards compatibility
ffn_layer = "dense_relu_dense"
if ffn_layer == "dense_relu_dense":
# In simple convolution mode, use `pad_remover` to speed up processing.
mlperf_log.transformer_print(
key=mlperf_log.MODEL_HP_FFN_FILTER_DENSE,
value={
"filter_size": hparams.filter_size,
"use_bias": "True",
"activation": mlperf_log.RELU
})
mlperf_log.transformer_print(
key=mlperf_log.MODEL_HP_FFN_OUTPUT_DENSE,
value={
"hidden_size": hparams.hidden_size,
"use_bias": "True",
})
mlperf_log.transformer_print(
key=mlperf_log.MODEL_HP_RELU_DROPOUT, value=hparams.relu_dropout)
if pad_remover:
original_shape = common_layers.shape_list(x)
# Collapse `x` across examples, and remove padding positions.
x = tf.reshape(x, tf.concat([[-1], original_shape[2:]], axis=0))
x = tf.expand_dims(pad_remover.remove(x), axis=0)
conv_output = common_layers.dense_relu_dense(
x,
hparams.filter_size,
hparams.hidden_size,
dropout=hparams.relu_dropout,
dropout_broadcast_dims=relu_dropout_broadcast_dims)
if pad_remover:
# Restore `conv_output` to the original shape of `x`, including padding.
conv_output = tf.reshape(
pad_remover.restore(tf.squeeze(conv_output, axis=0)), original_shape)
return conv_output
elif ffn_layer == "conv_relu_conv":
return common_layers.conv_relu_conv(
x,
readout_filter_size or hparams.filter_size,
hparams.hidden_size,
first_kernel_size=hparams.conv_first_kernel,
second_kernel_size=1,
padding=conv_padding,
nonpadding_mask=nonpadding_mask,
dropout=hparams.relu_dropout,
cache=cache,
decode_loop_step=decode_loop_step)
elif ffn_layer == "parameter_attention":
return common_attention.parameter_attention(
x, hparams.parameter_attention_key_channels or hparams.hidden_size,
hparams.parameter_attention_value_channels or hparams.hidden_size,
hparams.hidden_size, readout_filter_size or hparams.filter_size,
hparams.num_heads,
hparams.attention_dropout)
elif ffn_layer == "conv_hidden_relu_with_sepconv":
return common_layers.conv_hidden_relu(
x,
readout_filter_size or hparams.filter_size,
hparams.hidden_size,
kernel_size=(3, 1),
second_kernel_size=(31, 1),
padding="LEFT",
dropout=hparams.relu_dropout)
elif ffn_layer == "sru":
return common_layers.sru(x)
elif ffn_layer == "local_moe_tpu":
overhead = (
hparams.moe_overhead_train
if hparams.mode == tf.estimator.ModeKeys.TRAIN else
hparams.moe_overhead_eval)
ret, loss = expert_utils.local_moe_tpu(
x,
hparams.filter_size // 2,
hparams.hidden_size,
hparams.moe_num_experts,
overhead=overhead,
loss_coef=hparams.moe_loss_coef)
elif ffn_layer == "local_moe":
overhead = (
hparams.moe_overhead_train
if hparams.mode == tf.estimator.ModeKeys.TRAIN else
hparams.moe_overhead_eval)
ret, loss = expert_utils.local_moe(
x,
True,
expert_utils.ffn_expert_fn(hparams.hidden_size, [hparams.filter_size],
hparams.hidden_size),
hparams.moe_num_experts,
k=hparams.moe_k,
hparams=hparams)
losses.append(loss)
return ret
else:
assert ffn_layer == "none"
return x
