def prepare_image_question_encoder(image_feat, question, hparams):
"""Prepare encoder.
Args:
image_feat: a Tensor.
question: a Tensor.
hparams: run hyperparameters
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
"""
encoder_input = tf.concat([image_feat, question], axis=1)
encoder_padding = common_attention.embedding_to_padding(encoder_input)
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
encoder_self_attention_bias = ignore_padding
encoder_decoder_attention_bias = ignore_padding
# Usual case - not a packed dataset.
if hparams.pos == "timing":
question = common_attention.add_timing_signal_1d(question)
elif hparams.pos == "emb":
question = common_attention.add_positional_embedding(
question, hparams.max_length, "inputs_positional_embedding",
None)
encoder_input = tf.concat([image_feat, question], axis=1)
return (encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias)
def prepare_image_question_encoder(image_feat, question, hparams):
"""Prepare encoder.
Args:
image_feat: a Tensor.
question: a Tensor.
hparams: run hyperparameters
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
"""
encoder_input = tf.concat([image_feat, question], axis=1)
encoder_padding = common_attention.embedding_to_padding(encoder_input)
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
encoder_self_attention_bias = ignore_padding
encoder_decoder_attention_bias = ignore_padding
# Usual case - not a packed dataset.
if hparams.pos == "timing":
question = common_attention.add_timing_signal_1d(question)
elif hparams.pos == "emb":
question = common_attention.add_positional_embedding(
question, hparams.max_length, "inputs_positional_embedding", None)
encoder_input = tf.concat([image_feat, question], axis=1)
return (encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias)
def testAddPositionalEmbedding(self):
x = np.random.rand(5, 3, 12)
y = common_attention.add_positional_embedding(
tf.constant(x, dtype=tf.float32),
max_length=4,
name="pos_embedding")
self.evaluate(tf.global_variables_initializer())
res = self.evaluate(y)
self.assertEqual(res.shape, (5, 3, 12))
def transformer_prepare_decoder_right(targets, hparams, features=None):
"""Prepare one shard of the model for the decoder.
Args:
targets: a Tensor.
hparams: run hyperparameters
features: optionally pass the entire features dictionary as well.
This is needed now for "packed" datasets.
Returns:
decoder_input: a Tensor, bottom of decoder stack
decoder_self_attention_bias: a bias tensor for use in decoder self-attention
"""
if hparams.causal_decoder_self_attention:
# Causal attention.
if hparams.prepend_mode == "prepend_inputs_full_attention":
decoder_self_attention_bias = (
common_attention.attention_bias_prepend_inputs_full_attention(
common_attention.embedding_to_padding(targets)))
else:
decoder_self_attention_bias = (
common_attention.attention_bias_local(
common_layers.shape_list(targets)[1], 0, -1))
else:
# Full attention.
decoder_padding = common_attention.embedding_to_padding(targets)
decoder_self_attention_bias = (
common_attention.attention_bias_ignore_padding(decoder_padding))
if features and "targets_segmentation" in features:
# "Packed" dataset - keep the examples from seeing each other.
targets_segmentation = features["targets_segmentation"]
targets_position = features["targets_position"]
decoder_self_attention_bias += common_attention.attention_bias_same_segment(
targets_segmentation, targets_segmentation)
else:
targets_position = None
if hparams.proximity_bias:
decoder_self_attention_bias += common_attention.attention_bias_proximal(
common_layers.shape_list(targets)[1])
decoder_input = shift_left_3d(targets)
if hparams.pos == "timing":
if targets_position is not None:
decoder_input = common_attention.add_timing_signal_1d_given_position(
decoder_input, targets_position)
else:
decoder_input = common_attention.add_timing_signal_1d(
decoder_input)
elif hparams.pos == "emb":
decoder_input = common_attention.add_positional_embedding(
decoder_input, hparams.max_length, "targets_positional_embedding",
targets_position)
if hparams.activation_dtype == "bfloat16":
decoder_self_attention_bias = tf.cast(decoder_self_attention_bias,
tf.bfloat16)
return (decoder_input, decoder_self_attention_bias)
def transformer_prepare_encoder(inputs_emb_var,
inputs,
hparams,
features=None):
"""Prepare one shard of the model for the encoder.
Args:
inputs_emb_var: a Tensor
inputs: a Tensor.
hparams: run hyperparameters
features: optionally pass the entire features dictionary as well.
This is needed now for "packed" datasets.
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
attention
"""
encoder_input = tf.gather(inputs_emb_var, inputs)
if features and "inputs_segmentation" in features:
# Packed dataset. Keep the examples from seeing each other.
inputs_segmentation = features["inputs_segmentation"]
inputs_position = features["inputs_position"]
targets_segmentation = features["targets_segmentation"]
encoder_self_attention_bias = common_attention.attention_bias_same_segment(
inputs_segmentation, inputs_segmentation)
encoder_decoder_attention_bias = (
common_attention.attention_bias_same_segment(
targets_segmentation, inputs_segmentation))
else:
# Usual case - not a packed dataset.
encoder_padding = tf.to_float(tf.equal(inputs, 0))
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
encoder_self_attention_bias = ignore_padding
encoder_decoder_attention_bias = ignore_padding
inputs_position = None
with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
encoder_input = common_attention.add_positional_embedding(
encoder_input, hparams.max_length, "positional_embedding",
inputs_position)
if hparams.activation_dtype == "bfloat16":
encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
tf.bfloat16)
encoder_decoder_attention_bias = tf.cast(
encoder_decoder_attention_bias, tf.bfloat16)
return (encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias)
def transformer_prepare_decoder(targets_emb_var,
targets,
hparams,
features=None):
"""Prepare one shard of the model for the decoder.
Args:
targets_emb_var: a Tensor
targets: a Tensor.
hparams: run hyperparameters
features: optionally pass the entire features dictionary as well.
This is needed now for "packed" datasets.
Returns:
decoder_input: a Tensor, bottom of decoder stack
decoder_self_attention_bias: a bias tensor for use in decoder self-attention
"""
decoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(
common_layers.shape_list(targets)[1]))
if features and "targets_segmentation" in features:
# "Packed" dataset - keep the examples from seeing each other.
targets_segmentation = features["targets_segmentation"]
targets_position = features["targets_position"]
decoder_self_attention_bias += common_attention.attention_bias_same_segment(
targets_segmentation, targets_segmentation)
else:
targets_position = None
decoder_input = tf.gather(targets_emb_var,
common_layers.shift_right_2d(targets))
with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
decoder_input = common_attention.add_positional_embedding(
decoder_input, hparams.max_length, "positional_embedding",
targets_position)
if hparams.activation_dtype == "bfloat16":
decoder_self_attention_bias = tf.cast(decoder_self_attention_bias,
tf.bfloat16)
return (decoder_input, decoder_self_attention_bias)
def prepare_question_encoder(inputs, hparams):
"""Prepare question encoder.
Args:
inputs: a Tensor.
hparams: run hyperparameters
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
"""
encoder_input = inputs
# Usual case - not a packed dataset.
encoder_padding = common_attention.embedding_to_padding(encoder_input)
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
encoder_self_attention_bias = ignore_padding
if hparams.pos == "timing":
encoder_input = common_attention.add_timing_signal_1d(encoder_input)
elif hparams.pos == "emb":
encoder_input = common_attention.add_positional_embedding(
encoder_input, hparams.max_length, "inputs_positional_embedding",
None)
return (encoder_input, encoder_self_attention_bias)
def prepare_question_encoder(inputs, hparams):
"""Prepare question encoder.
Args:
inputs: a Tensor.
hparams: run hyperparameters
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
"""
encoder_input = inputs
# Usual case - not a packed dataset.
encoder_padding = common_attention.embedding_to_padding(encoder_input)
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
encoder_self_attention_bias = ignore_padding
if hparams.pos == "timing":
encoder_input = common_attention.add_timing_signal_1d(encoder_input)
elif hparams.pos == "emb":
encoder_input = common_attention.add_positional_embedding(
encoder_input, hparams.max_length, "inputs_positional_embedding",
None)
return (encoder_input, encoder_self_attention_bias)
def _fast_decode(self,
features,
decode_length,
beam_size=1,
top_beams=1,
alpha=1.0):
if self._num_datashards != 1:
raise NotImplementedError(
"Fast decoding only supports a single shard.")
dp = self._data_parallelism
hparams = self._hparams
target_modality = self._problem_hparams.modality["targets"]
if "targets_segmentation" in features:
raise NotImplementedError(
"Decoding not supported on packed datasets "
" If you want to decode from a dataset, use the non-packed version"
" of the dataset when decoding.")
if self.has_input:
inputs = features["inputs"]
if target_modality.is_class_modality:
decode_length = 1
else:
decode_length = (common_layers.shape_list(inputs)[1] +
features.get("decode_length", decode_length))
contexts = {}
for feature_name in features:
if 'context' in feature_name and 'raw' not in feature_name:
contexts[feature_name] = features[feature_name]
inputs = tf.expand_dims(inputs, axis=1)
if len(inputs.shape) < 5:
inputs = tf.expand_dims(inputs, axis=4)
s = common_layers.shape_list(inputs)
batch_size = s[0]
inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
# _shard_features called to ensure that the variable names match
inputs = self._shard_features({"inputs": inputs})["inputs"]
input_modality = self._problem_hparams.modality["inputs"]
context_modality = {}
for context_name in contexts:
if context_name in self._problem_hparams.modality:
context_modality[
context_name] = self._problem_hparams.modality[
context_name]
else:
context_modality[context_name] = input_modality
with tf.variable_scope(input_modality.name, reuse=tf.AUTO_REUSE):
inputs = input_modality.bottom_sharded(inputs, dp)
for feature_name in contexts:
with tf.variable_scope(context_modality[feature_name].name,
reuse=tf.AUTO_REUSE):
contexts[feature_name] = context_modality[
feature_name].bottom_sharded(contexts[feature_name],
dp)
contexts_list = [
contexts[feature_name][0] for feature_name in contexts
]
contexts = tf.concat(contexts_list, axis=1)
inputs = [tf.concat([contexts, inputs[0]], axis=1)]
with tf.variable_scope("body"):
encoder_output, encoder_decoder_attention_bias = dp(
self.encode,
inputs,
features["target_space_id"],
hparams,
features=features)
encoder_output = encoder_output[0]
encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
partial_targets = None
else:
# The problem has no inputs.
encoder_output = None
encoder_decoder_attention_bias = None
# Prepare partial targets.
# In either features["inputs"] or features["targets"].
# We force the outputs to begin with these sequences.
partial_targets = features.get("inputs")
if partial_targets is None:
partial_targets = features["targets"]
assert partial_targets is not None
partial_targets = common_layers.expand_squeeze_to_nd(
partial_targets, 2)
partial_targets = tf.to_int64(partial_targets)
partial_targets_shape = common_layers.shape_list(partial_targets)
partial_targets_length = partial_targets_shape[1]
decode_length = (partial_targets_length +
features.get("decode_length", decode_length))
batch_size = partial_targets_shape[0]
if hparams.pos == "timing":
positional_encoding = common_attention.get_timing_signal_1d(
decode_length + 1, hparams.hidden_size)
elif hparams.pos == "emb":
positional_encoding = common_attention.add_positional_embedding(
tf.zeros([1, decode_length, hparams.hidden_size]),
hparams.max_length, "body/targets_positional_embedding", None)
else:
positional_encoding = None
def preprocess_targets(targets, i):
"""Performs preprocessing steps on the targets to prepare for the decoder.
This includes:
- Embedding the ids.
- Flattening to 3D tensor.
- Optionally adding timing signals.
Args:
targets: inputs ids to the decoder. [batch_size, 1]
i: scalar, Step number of the decoding loop.
Returns:
Processed targets [batch_size, 1, hidden_dim]
"""
# _shard_features called to ensure that the variable names match
targets = self._shard_features({"targets": targets})["targets"]
with tf.variable_scope(target_modality.name):
targets = target_modality.targets_bottom_sharded(targets,
dp)[0]
targets = common_layers.flatten4d3d(targets)
targets = tf.cond(tf.equal(i, 0), lambda: tf.zeros_like(targets),
lambda: targets)
if positional_encoding is not None:
targets += positional_encoding[:, i:i + 1]
return targets
decoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(decode_length))
if hparams.proximity_bias:
decoder_self_attention_bias += common_attention.attention_bias_proximal(
decode_length)
def symbols_to_logits_fn(ids, i, cache):
"""Go from ids to logits for next symbol."""
ids = ids[:, -1:]
targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
targets = preprocess_targets(targets, i)
bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
with tf.variable_scope("body"):
body_outputs = dp(self.decode,
targets,
cache.get("encoder_output"),
cache.get("encoder_decoder_attention_bias"),
bias,
hparams,
cache,
nonpadding=features_to_nonpadding(
features, "targets"))
with tf.variable_scope(target_modality.name):
logits = target_modality.top_sharded(body_outputs, None, dp)[0]
ret = tf.squeeze(logits, axis=[1, 2, 3])
if partial_targets is not None:
# If the position is within the given partial targets, we alter the
# logits to always return those values.
# A faster approach would be to process the partial targets in one
# iteration in order to fill the corresponding parts of the cache.
# This would require broader changes, though.
vocab_size = tf.shape(ret)[1]
def forced_logits():
return tf.one_hot(
tf.tile(partial_targets[:, i], [beam_size]),
vocab_size, 0.0, -1e9)
ret = tf.cond(tf.less(i, partial_targets_length),
forced_logits, lambda: ret)
return ret, cache
ret = fast_decode(
encoder_output=encoder_output,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
symbols_to_logits_fn=symbols_to_logits_fn,
hparams=hparams,
decode_length=decode_length,
vocab_size=target_modality.top_dimensionality,
beam_size=beam_size,
top_beams=top_beams,
alpha=alpha,
batch_size=batch_size,
force_decode_length=self._decode_hparams.force_decode_length)
if partial_targets is not None:
if beam_size <= 1 or top_beams <= 1:
ret["outputs"] = ret["outputs"][:, partial_targets_length:]
else:
ret["outputs"] = ret["outputs"][:, :, partial_targets_length:]
return ret
def _fast_decode(self,
features,
decode_length,
beam_size=1,
top_beams=1,
alpha=1.0):
#dp = self._data_parallelism
hparams = self._hparams
target_modality = self._problem_hparams.modality["targets"]
inputs = features["inputs"]
decode_length = (common_layers.shape_list(inputs)[1] +
features.get("decode_length", decode_length))
#inputs = tf.expand_dims(inputs, axis=1)
#if len(inputs.shape) < 5:
# inputs = tf.expand_dims(inputs, axis=4)
s = common_layers.shape_list(inputs)
batch_size = s[0]
#inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
# _shard_features called to ensure that the variable names match
#inputs = self._shard_features({"inputs": inputs})["inputs"]
input_modality = self._problem_hparams.modality["inputs"]
context_modality = {}
contexts = {}
for feature_name in features:
if 'context' in feature_name and 'raw' not in feature_name:
contexts[feature_name] = features[feature_name]
for context_name in contexts:
if context_name in self._problem_hparams.modality:
context_modality[
context_name] = self._problem_hparams.modality[
context_name]
else:
context_modality[context_name] = input_modality
with tf.variable_scope(input_modality.name, reuse=tf.AUTO_REUSE):
inputs = input_modality.bottom(inputs)
for context_name in contexts:
contexts[context_name] = context_modality[context_name].bottom(
contexts[context_name])
with tf.variable_scope("body", reuse=tf.AUTO_REUSE):
encoder_output, encoder_decoder_attention_bias = self.encode(
inputs,
contexts,
features["target_space_id"],
hparams,
features=features)
#encoder_output = encoder_output[0]
#encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
partial_targets = None
if hparams.pos == "timing":
positional_encoding = common_attention.get_timing_signal_1d(
decode_length + 1, hparams.hidden_size)
elif hparams.pos == "emb":
positional_encoding = common_attention.add_positional_embedding(
tf.zeros([1, decode_length + 1, hparams.hidden_size]),
hparams.max_length, "targets_positional_embedding", None)
else:
positional_encoding = None
def preprocess_targets(targets, i):
"""Performs preprocessing steps on the targets to prepare for the decoder.
This includes:
- Embedding the ids.
- Flattening to 3D tensor.
- Optionally adding timing signals.
Args:
targets: inputs ids to the decoder. [batch_size, 1]
i: scalar, Step number of the decoding loop.
Returns:
Processed targets [batch_size, 1, hidden_dim]
"""
# _shard_features called to ensure that the variable names match
#targets = self._shard_features({"targets": targets})["targets"]
with tf.variable_scope(target_modality.name):
targets = target_modality.targets_bottom(targets)
targets = common_layers.flatten4d3d(targets)
targets = tf.cond(tf.equal(i, 0), lambda: tf.zeros_like(targets),
lambda: targets)
if positional_encoding is not None:
targets += positional_encoding[:, i:i + 1]
return targets
decoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(decode_length))
if hparams.proximity_bias:
decoder_self_attention_bias += common_attention.attention_bias_proximal(
decode_length)
def symbols_to_logits_fn(ids, i, cache):
"""Go from ids to logits for next symbol."""
ids = ids[:, -1:]
targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
targets = preprocess_targets(targets, i)
bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
with tf.variable_scope("body"):
body_outputs = self.decode(
targets,
cache.get("encoder_output"),
cache.get("encoder_decoder_attention_bias"),
bias,
hparams,
cache,
nonpadding=features_to_nonpadding(features, "targets"))
with tf.variable_scope(target_modality.name):
logits = target_modality.top(body_outputs, None)
ret = tf.squeeze(logits, axis=[1, 2, 3])
return ret, cache
ret = fast_decode(
encoder_output=encoder_output,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
symbols_to_logits_fn=symbols_to_logits_fn,
hparams=hparams,
decode_length=decode_length,
vocab_size=target_modality.top_dimensionality,
beam_size=beam_size,
top_beams=top_beams,
alpha=alpha,
batch_size=batch_size,
force_decode_length=self._decode_hparams.force_decode_length)
return ret
def ae_transformer_internal(inputs,
targets,
target_space,
hparams,
cache=None,
predict_mask=1.0):
"""AE Transformer, main step used for training."""
# Summaries break with the do_refine cond, turn them off in that case.
global _DO_SUMMARIES
if hparams.do_refine:
_DO_SUMMARIES = False
# Prepare.
if inputs is not None:
batch_size = common_layers.shape_list(inputs)[0]
else:
batch_size = common_layers.shape_list(targets)[0]
targets = tf.reshape(targets, [batch_size, -1, 1, hparams.hidden_size])
# Encoder.
if inputs is not None:
inputs = common_layers.flatten4d3d(inputs)
inputs, ed = encode(inputs, target_space, hparams, "input_enc")
inputs_ex, ed_ex = inputs, ed
else:
ed, inputs_ex, ed_ex = None, None, None
# Autoencoding.
losses = {
"extra": tf.constant(0.0),
"latent_pred": tf.constant(0.0),
"neg_q_entropy": tf.constant(0.0)
}
if hparams.do_ae:
# flatten here
original_targets = targets
original_targets_shape = tf.shape(original_targets)
if hparams.task == "image":
cia.maybe_reshape_4d_to_3d(targets)
if hparams.task == "translate":
if inputs is not None:
max_targets_len_from_inputs = tf.concat([inputs, inputs],
axis=1)
else:
max_targets_len_from_inputs = targets
else:
assert hparams.task == "image"
max_targets_len_from_inputs = targets
if hparams.word_shuffle:
tf.logging.info("Using word shuffle with rate = {}".format(
hparams.word_shuffle))
targets_idx = tf.range(start=0,
limit=common_layers.shape_list(targets)[1],
delta=1)
targets_idx = tf.to_float(targets_idx)
noise = tf.random_uniform(
shape=common_layers.shape_list(targets_idx),
minval=0,
maxval=1 + hparams.word_shuffle)
targets_idx += noise
permutation = contrib.framework().argsort(targets_idx)
targets_permuted = tf.gather(targets, indices=permutation, axis=1)
targets = targets_permuted
targets, _ = common_layers.pad_to_same_length(
targets,
max_targets_len_from_inputs,
final_length_divisible_by=2**hparams.num_compress_steps)
# Add positional information
targets_shape = common_layers.shape_list(targets)
targets = tf.reshape(
targets, [targets_shape[0], targets_shape[1], targets_shape[3]])
targets = common_attention.add_positional_embedding(
targets, hparams.max_length, name="targets_position")
targets = tf.reshape(targets, shape=targets_shape)
if hparams.word_dropout:
mask = tf.random_uniform(shape=common_layers.shape_list(targets),
minval=0.0,
maxval=1.0)
targets_noisy = tf.where(mask > hparams.word_dropout, targets,
tf.zeros_like(targets))
else:
targets_noisy = targets
targets_c = compress(targets_noisy, inputs, False, hparams, "compress")
if hparams.mode != tf.estimator.ModeKeys.PREDICT:
# Compress and bottleneck.
latents_dense, latents_discrete, extra_loss, embed, neg_q_entropy = (
hparams.bottleneck(inputs=targets_c,
filter_size=hparams.compress_filter_size,
mode=hparams.mode,
name="vc"))
if _DO_SUMMARIES:
tf.summary.histogram(
"b0", tf.reshape(latents_discrete[:, 0, :], [-1]))
pc = common_layers.inverse_exp_decay(hparams.startup_steps)
pc = pc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
cond = tf.less(tf.random_uniform([batch_size]), pc)
latents_dense = tf.where(cond, latents_dense, targets_c)
# TODO(lukaszkaiser): return extra losses batchwise, multiply before mean.
losses["extra"] = extra_loss * tf.reduce_mean(tf.to_float(cond))
# Extra loss predicting latent code from input. Discrete only.
if hparams.bottleneck_kind not in ["dense", "vae"]:
latents_pred = decode_transformer(inputs_ex,
ed_ex,
embed(latents_discrete),
hparams,
"extra",
task="translate")
_, latent_pred_loss = ae_latent_softmax(
latents_pred, tf.stop_gradient(latents_discrete), hparams)
# Scale by latent dimension for summary so we can compare across
# batches.
if _DO_SUMMARIES:
tf.summary.scalar("latent_pred_loss_mean",
tf.reduce_mean(latent_pred_loss))
if hparams.sum_over_latents:
latent_pred_loss = tf.reduce_sum(latent_pred_loss, [1, 2])
losses["latent_pred"] = tf.reduce_mean(
latent_pred_loss * tf.to_float(cond)) * hparams.prior_scale
losses["neg_q_entropy"] = neg_q_entropy * hparams.entropy_scale
else:
inputs_c = decode_transformer(inputs, ed, targets_c, hparams,
"dec_c")
losses["latent_pred"] = tf.reduce_mean(
tf.squared_difference(inputs_c, targets_c)) * 20
def bn_inputs():
with tf.variable_scope(tf.get_variable_scope(),
reuse=True):
bn, _, _, _, _ = hparams.bottleneck(
inputs=inputs_c,
filter_size=hparams.compress_filter_size,
mode=hparams.mode,
name="vc")
return bn
inputs_c = bn_inputs()
ptc = 1.0 - common_layers.inverse_lin_decay(200000) * 0.5
ptc = ptc if hparams.mode == tf.estimator.ModeKeys.TRAIN else 1.0
latents_dense = tf.where(
tf.less(tf.random_uniform([batch_size]), ptc),
latents_dense, inputs_c)
else:
if hparams.bottleneck_kind in ["dense", "vae"]:
inputs_c = decode_transformer(inputs, ed, targets_c, hparams,
"dec_c")
latents_dense, _, _, _, _ = hparams.bottleneck(
inputs=inputs_c,
filter_size=hparams.compress_filter_size,
mode=hparams.mode,
name="vc")
else:
latent_len = common_layers.shape_list(targets_c)[1]
_, _, _, embed, _ = hparams.bottleneck(
inputs=targets_c,
filter_size=hparams.compress_filter_size,
name="vc")
latents_dense = tf.zeros_like(targets_c[:, :latent_len, :, :])
if cache is None:
cache = ae_latent_sample(latents_dense, inputs_ex, ed_ex,
embed, 16, hparams)
latents_dense = embed(cache)
# Postprocess.
d = latents_dense
d_shape = common_layers.shape_list(d)
d = tf.reshape(d, [d_shape[0], d_shape[1], d_shape[3]])
d = common_attention.add_positional_embedding(d,
hparams.max_length,
name="latents_position")
d = tf.reshape(d, shape=d_shape)
# decompressing the dense latents
for i in range(hparams.num_compress_steps):
j = hparams.num_compress_steps - i - 1
d = residual_conv(d, 1, (3, 1), hparams, "decompress_rc_%d" % j)
if inputs is not None and hparams.do_attend_decompress:
d = attend(d, inputs, hparams, "decompress_attend_%d" % j)
d = decompress_step(d, hparams, i > 0, False, "decompress_%d" % j)
# Masking.
if hparams.do_mask:
masking = common_layers.inverse_lin_decay(
hparams.mask_startup_steps)
masking *= common_layers.inverse_exp_decay(
hparams.mask_startup_steps // 4) # Not much at start.
if not hparams.do_refine:
masking -= tf.random_uniform([]) * hparams.unmasked_percentage
masking = tf.minimum(tf.maximum(masking, 0.0), 1.0)
if hparams.use_predict_mask:
masking = predict_mask
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
masking = predict_mask
mask = tf.less(
masking,
tf.random_uniform(common_layers.shape_list(targets)[:-1]))
mask = tf.expand_dims(tf.to_float(mask), 3)
# targets is always [batch, length, 1, depth]
targets = mask * targets + (1.0 - mask) * d
# reshape back to 4d here
if hparams.task == "image":
targets = tf.reshape(targets, original_targets_shape)
else:
targets = d
res = decode_transformer(inputs,
ed,
targets,
hparams,
"decoder",
causal=hparams.causal)
if hparams.do_ae:
if hparams.do_mask and hparams.do_refine:
def refine_res():
# return residual_conv(res, 1, (5, 1), hparams, "refine")
r, _ = encode(tf.squeeze(res, axis=[2]), target_space, hparams,
"refine_enc")
return tf.expand_dims(r, axis=2)
masked_batches = tf.reduce_sum(mask, axis=[1, 2, 3])
all_masked = tf.less(masked_batches, 0.1)
res = tf.where(all_masked, refine_res(), res)
# We'll start training the extra model of latents after mask_startup_steps.
nonlatent_steps = hparams.mask_startup_steps
latent_time = tf.less(nonlatent_steps,
tf.to_int32(tf.train.get_global_step()))
losses["latent_pred"] *= tf.to_float(latent_time)
# res was generated from padded targets, which means it has some extra
# elements. These can cause shape problems when computing loss with respect to
# the original (unpadded) targets. So we remove their extra elements here.
res = res[:, :original_targets_shape[1], :, :]
data_dim = common_layers.shape_list(res)[1]
latent_dim = common_layers.shape_list(targets_c)[1]
return res, losses, cache, data_dim, latent_dim
def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
"""Prepare one shard of the model for the encoder.
Args:
inputs: a Tensor.
target_space: a Tensor.
hparams: run hyperparameters
features: optionally pass the entire features dictionary as well.
This is needed now for "packed" datasets.
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
attention
"""
ishape_static = inputs.shape.as_list()
encoder_input = inputs
if features and "inputs_segmentation" in features:
# Packed dataset. Keep the examples from seeing each other.
inputs_segmentation = features["inputs_segmentation"]
inputs_position = features["inputs_position"]
targets_segmentation = features["targets_segmentation"]
if (hasattr(hparams, "unidirectional_encoder") and
hparams.unidirectional_encoder):
tf.logging.info("Using unidirectional encoder")
encoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(
common_layers.shape_list(inputs)[1]))
else:
encoder_self_attention_bias = (
common_attention.attention_bias_same_segment(
inputs_segmentation, inputs_segmentation))
encoder_decoder_attention_bias = (
common_attention.attention_bias_same_segment(targets_segmentation,
inputs_segmentation))
else:
encoder_padding = common_attention.embedding_to_padding(encoder_input)
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
if (hasattr(hparams, "unidirectional_encoder") and
hparams.unidirectional_encoder):
tf.logging.info("Using unidirectional encoder")
encoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(
common_layers.shape_list(inputs)[1]))
else:
# Usual case - not a packed dataset.
encoder_self_attention_bias = ignore_padding
encoder_decoder_attention_bias = ignore_padding
inputs_position = None
if hparams.proximity_bias:
encoder_self_attention_bias += common_attention.attention_bias_proximal(
common_layers.shape_list(inputs)[1])
if hparams.get("use_target_space_embedding", True):
# Append target_space_id embedding to inputs.
emb_target_space = common_layers.embedding(
target_space,
32,
ishape_static[-1],
name="target_space_embedding",
dtype=tf.bfloat16
if hparams.activation_dtype == "bfloat16" else tf.float32)
emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
encoder_input += emb_target_space
if hparams.pos == "timing":
if inputs_position is not None:
encoder_input = common_attention.add_timing_signal_1d_given_position(
encoder_input, inputs_position)
else:
encoder_input = common_attention.add_timing_signal_1d(encoder_input)
elif hparams.pos == "emb":
encoder_input = common_attention.add_positional_embedding(
encoder_input, hparams.max_length, "inputs_positional_embedding",
inputs_position)
if hparams.activation_dtype == "bfloat16":
encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
tf.bfloat16)
encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
tf.bfloat16)
return (encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias)
def _fast_decode(
self,
features,
decode_length,
beam_size=1,
top_beams=1,
alpha=1.0,
preprocess_targets_method=None,
):
if self._num_datashards != 1:
raise NotImplementedError(
'Fast decoding only supports a single shard.')
dp = self._data_parallelism
hparams = self._hparams
target_modality = self._problem_hparams.modality['targets']
target_vocab_size = self._problem_hparams.vocab_size['targets']
if target_vocab_size is not None and hasattr(hparams, 'vocab_divisor'):
target_vocab_size += (-target_vocab_size) % hparams.vocab_divisor
target_tag_modality = self._problem_hparams.modality[
'targets_error_tag']
target_tag_vocab_size = self._problem_hparams.vocab_size[
'targets_error_tag']
if target_tag_vocab_size is not None and hasattr(
hparams, 'vocab_divisor'):
target_tag_vocab_size += (
-target_tag_vocab_size) % hparams.vocab_divisor
if 'targets_segmentation' in features:
raise NotImplementedError(
'Decoding not supported on packed datasets '
' If you want to decode from a dataset, use the non-packed version'
' of the dataset when decoding.')
if self.has_input:
inputs_shape = common_layers.shape_list(features['inputs'])
if (target_modality == modalities.ModalityType.CLASS_LABEL
or self._problem_hparams.get('regression_targets')):
decode_length = 1
else:
decode_length = inputs_shape[1] + features.get(
'decode_length', decode_length)
batch_size = inputs_shape[0]
inputs = self._prepare_inputs_for_decode(features)
with tf.variable_scope('body'):
encoder_output, encoder_decoder_attention_bias = dp(
self.encode,
inputs,
features['target_space_id'],
hparams,
features=features,
)
encoder_output = encoder_output[0]
encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
partial_targets = features.get('partial_targets')
else:
encoder_output = None
encoder_decoder_attention_bias = None
partial_targets = features.get('inputs')
if partial_targets is None:
partial_targets = features['targets']
assert partial_targets is not None
if partial_targets is not None:
partial_targets = common_layers.expand_squeeze_to_nd(
partial_targets, 2)
partial_targets = tf.to_int64(partial_targets)
partial_targets_shape = common_layers.shape_list(partial_targets)
partial_targets_length = partial_targets_shape[1]
decode_length = partial_targets_length + features.get(
'decode_length', decode_length)
batch_size = partial_targets_shape[0]
if hparams.pos == 'timing':
positional_encoding = common_attention.get_timing_signal_1d(
decode_length + 1, hparams.hidden_size)
elif hparams.pos == 'timing_from_features':
positional_encoding = common_attention.add_timing_signals_from_features(
tf.zeros([1, decode_length, hparams.hidden_size]),
features,
hparams.position_features,
)
elif hparams.pos == 'emb':
positional_encoding = common_attention.add_positional_embedding(
tf.zeros([1, decode_length, hparams.hidden_size]),
hparams.max_length,
'body/targets_positional_embedding',
None,
)
else:
positional_encoding = None
def preprocess_targets(targets, i):
targets = self._shard_features({'targets': targets})['targets']
modality_name = hparams.name.get(
'targets',
modalities.get_name(target_modality))(hparams,
target_vocab_size)
with tf.variable_scope(modality_name + '/targets'):
bottom = hparams.bottom.get(
'targets', modalities.get_targets_bottom(target_modality))
targets = dp(bottom, targets, hparams, target_vocab_size)[0]
targets = common_layers.flatten4d3d(targets)
if not self.get_decode_start_id():
targets = tf.cond(
tf.equal(i, 0),
lambda: tf.zeros_like(targets),
lambda: targets,
)
if positional_encoding is not None:
targets += positional_encoding[:, i:i + 1]
return targets
def preprocess_targets_tag_method(targets, i):
targets = self._shard_features({'targets_error_tag':
targets})['targets_error_tag']
modality_name = hparams.name.get(
'targets_error_tag', modalities.get_name(target_tag_modality))(
hparams, target_tag_vocab_size)
with tf.variable_scope(modality_name + '/targets_error_tag'):
bottom = hparams.bottom.get(
'targets_error_tag',
modalities.get_targets_bottom(target_tag_modality),
)
targets = dp(bottom, targets, hparams,
target_tag_vocab_size)[0]
targets = common_layers.flatten4d3d(targets)
if not self.get_decode_start_id():
targets = tf.cond(
tf.equal(i, 0),
lambda: tf.zeros_like(targets),
lambda: targets,
)
if positional_encoding is not None:
targets += positional_encoding[:, i:i + 1]
return targets
decoder_self_attention_bias = common_attention.attention_bias_lower_triangle(
decode_length)
if hparams.proximity_bias:
decoder_self_attention_bias += common_attention.attention_bias_proximal(
decode_length)
att_cache = {'attention_history': {}}
num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
if encoder_output is not None:
att_batch_size, enc_seq_length = common_layers.shape_list(
encoder_output)[0:2]
for layer in range(num_layers):
att_cache['attention_history']['layer_%d' % layer] = tf.zeros(
[att_batch_size, hparams.num_heads, 0, enc_seq_length])
def update_decoder_attention_history(cache):
for k in [
x for x in self.attention_weights
if 'decoder' in x and 'self' not in x and 'logits' not in x
]:
idx = k.find('layer_')
if idx < 0:
continue
# Get layer number from the string name.
layer_nbr = k[idx + 6:]
idx = 0
while (idx + 1 < len(layer_nbr)
and layer_nbr[:idx + 1].isdigit()):
idx += 1
layer_nbr = 'layer_%d' % int(layer_nbr[:idx])
if layer_nbr in cache['attention_history']:
cache['attention_history'][layer_nbr] = tf.concat(
[
cache['attention_history'][layer_nbr],
self.attention_weights[k],
],
axis=2,
)
if not preprocess_targets_method:
preprocess_targets_method = preprocess_targets
def symbols_to_logits_fn(ids, ids_tag, i, cache):
"""Go from ids to logits for next symbol."""
ids = ids[:, -1:]
targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
targets = preprocess_targets_method(targets, i)
ids_tag = ids_tag[:, -1:]
targets_tag = tf.expand_dims(tf.expand_dims(ids_tag, axis=2),
axis=3)
targets_tag = preprocess_targets_tag_method(targets_tag, i)
bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
with tf.variable_scope('body'):
with tf.variable_scope('edit_ops_layer'):
with tf.variable_scope('ffn'):
x = targets
preproc = lambda z: common_layers.layer_preprocess(
z, hparams, layer_collection=None)
layer_inputs = [
tf.concat(preproc(x), axis=0),
tf.concat(preproc(targets_tag), axis=0),
]
y = transformer_layers.transformer_ffn_layer(
tf.concat(layer_inputs, axis=2),
hparams,
conv_padding='LEFT',
nonpadding_mask=features_to_nonpadding(
features, 'targets'),
losses=None,
cache=cache,
decode_loop_step=None,
layer_collection=None,
)
targets = common_layers.layer_postprocess(
x, y, hparams)
if hparams.middle_prediction:
num_decoder_layers = (hparams.num_decoder_layers
or hparams.num_hidden_layers)
hparams.num_decoder_layers = int(
num_decoder_layers /
hparams.middle_prediction_layer_factor)
body_outputs = dp(
self.decode,
targets,
cache.get('encoder_output'),
cache.get('encoder_decoder_attention_bias'),
bias,
hparams,
cache,
nonpadding=features_to_nonpadding(features, 'targets'),
)[0]
body_outputs, logits_tag = dp(
self._prediction_cascade_predict,
hparams,
features_to_nonpadding(features, 'targets'),
cache.get('encoder_decoder_attention_bias'),
cache.get('encoder_output'),
body_outputs,
)
logits_tag = logits_tag[0]['targets_error_tag']
if hparams.middle_prediction:
with tf.variable_scope('after_prediction'):
body_outputs = dp(
self.decode,
targets + body_outputs[0],
cache.get('encoder_output'),
cache.get('encoder_decoder_attention_bias'),
bias,
hparams,
cache,
nonpadding=features_to_nonpadding(
features, 'targets'),
)
update_decoder_attention_history(cache)
modality_name = hparams.name.get(
'targets',
modalities.get_name(target_modality))(hparams,
target_vocab_size)
with tf.variable_scope('targets/' + modality_name):
top = hparams.top.get('targets',
modalities.get_top(target_modality))
logits = dp(top, body_outputs, None, hparams,
target_vocab_size)[0]
ret = tf.squeeze(logits, axis=[1, 2])
if partial_targets is not None:
vocab_size = tf.shape(ret)[1]
def forced_logits():
return tf.one_hot(
tf.tile(partial_targets[:, i], [beam_size]),
vocab_size,
0.0,
-1e9,
)
ret = tf.cond(
tf.less(i, partial_targets_length),
forced_logits,
lambda: ret,
)
logits_tag = tf.squeeze(logits_tag, axis=[1])
return ret, logits_tag, cache
sos_id = self.get_decode_start_id() or 0
eos_id = self.get_decode_end_id() or beam_search.EOS_ID
temperature = features.get('sampling_temp',
getattr(hparams, 'sampling_temp', 0.0))
top_k = features.get('sampling_keep_top_k',
getattr(hparams, 'sampling_keep_top_k', -1))
ret = fast_decode(
encoder_output=encoder_output,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
symbols_to_logits_fn=symbols_to_logits_fn,
hparams=hparams,
decode_length=decode_length,
vocab_size=target_vocab_size,
init_cache_fn=_init_transformer_cache,
beam_size=beam_size,
top_beams=top_beams,
alpha=alpha,
batch_size=batch_size,
force_decode_length=self._decode_hparams.force_decode_length,
sos_id=sos_id,
eos_id=eos_id,
sampling_temperature=temperature,
top_k=top_k,
cache=att_cache,
)
if partial_targets is not None:
if beam_size <= 1 or top_beams <= 1:
ret['outputs'] = ret['outputs'][:, partial_targets_length:]
else:
ret['outputs'] = ret['outputs'][:, :, partial_targets_length:]
return ret
def transformer_fn(self,
sentence_complex_input_placeholder, emb_complex,
sentence_simple_input_placeholder, emb_simple,
w, b,
rule_id_input_placeholder, rule_target_input_placeholder,
mem_contexts, mem_outputs,
global_step, score, comp_features, obj):
encoder_mask = tf.to_float(
tf.equal(tf.stack(sentence_complex_input_placeholder, axis=1),
self.data.vocab_complex.encode(constant.SYMBOL_PAD)))
encoder_attn_bias = common_attention.attention_bias_ignore_padding(encoder_mask)
obj_tensors = {}
train_mode = self.model_config.train_mode
if self.model_config.bert_mode:
# Leave space for decoder when static seq
gpu_id = 0 if train_mode == 'static_seq' or train_mode == 'static_self-critical' or 'direct' in self.model_config.memory else 1
with tf.device('/device:GPU:%s' % gpu_id):
sentence_complex_input = tf.stack(sentence_complex_input_placeholder, axis=1)
bert_model = BertModel(
BertConfig.from_json_file(self.model_config.bert_config),
self.is_train, sentence_complex_input,
input_mask=1.0-encoder_mask, token_type_ids=None, use_one_hot_embeddings=False)
encoder_embed_inputs = bert_model.embedding_output
encoder_outputs = bert_model.sequence_output
emb_complex = bert_model.embedding_table # update emb complex
if (self.model_config.tie_embedding == 'all' or
self.model_config.tie_embedding == 'enc_dec'):
emb_simple = bert_model.embedding_table
if (self.model_config.tie_embedding == 'all' or
self.model_config.tie_embedding == 'dec_out'):
emb_w_proj = tf.get_variable(
'emb_w_proj', shape=[self.model_config.dimension, self.model_config.dimension],
initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)
w = tf.matmul(bert_model.embedding_table, emb_w_proj)
if 'direct' in self.model_config.memory:
with tf.device('/device:GPU:1'):
direct_mask = tf.to_float(
tf.equal(tf.stack(rule_target_input_placeholder, axis=1),
self.data.vocab_complex.encode(constant.SYMBOL_PAD)))
direct_bert_model = BertModel(
BertConfig.from_json_file(self.model_config.bert_config),
self.is_train, tf.stack(rule_target_input_placeholder, axis=1),
input_mask=1.0 - direct_mask, token_type_ids=None, use_one_hot_embeddings=False,
embedding_table=emb_simple,
scope='direct')
direct_bert_output = direct_bert_model.sequence_output
obj_tensors['direct_bert_bias'] = common_attention.attention_bias_ignore_padding(direct_mask)
obj_tensors['direct_bert_output'] = direct_bert_output
else:
encoder_embed_inputs = tf.stack(
self.embedding_fn(sentence_complex_input_placeholder, emb_complex), axis=1)
if self.hparams.pos == 'timing':
encoder_embed_inputs = common_attention.add_timing_signal_1d(encoder_embed_inputs)
print('Use positional encoding in encoder text.')
if self.model_config.subword_vocab_size and self.model_config.seg_mode:
encoder_embed_inputs = common_attention.add_positional_embedding(
encoder_embed_inputs, 100, 'seg_embedding',
positions=obj['line_comp_segids'])
print('Add segment embedding.')
with tf.variable_scope('transformer_encoder'):
encoder_embed_inputs = tf.nn.dropout(encoder_embed_inputs,
1.0 - self.hparams.layer_prepostprocess_dropout)
if self.model_config.architecture == 'ut2t':
encoder_outputs, encoder_extra_output = universal_transformer_util.universal_transformer_encoder(
encoder_embed_inputs, encoder_attn_bias, self.hparams)
enc_ponder_times, enc_remainders = encoder_extra_output
extra_encoder_loss = (
self.hparams.act_loss_weight *
tf.reduce_mean(enc_ponder_times + enc_remainders))
if self.is_train:
obj_tensors['extra_encoder_loss'] = extra_encoder_loss
else:
encoder_outputs = transformer.transformer_encoder(
encoder_embed_inputs, encoder_attn_bias, self.hparams)
# Update score based on multiplier
score, pred_score_tuple = self.update_score(
score, encoder_outputs=encoder_outputs, encoder_mask=tf.to_float(
tf.not_equal(tf.stack(sentence_complex_input_placeholder, axis=1),
self.data.vocab_complex.encode(constant.SYMBOL_PAD))),
comp_features=comp_features)
encoder_outputs = self.update_encoder_embedding(encoder_outputs, score)
encoder_embed_inputs_list = tf.unstack(encoder_embed_inputs, axis=1)
with tf.variable_scope('transformer_decoder', reuse=tf.AUTO_REUSE):
if self.model_config.subword_vocab_size or 'bert_token' in self.model_config.bert_mode:
go_id = self.data.vocab_simple.encode(constant.SYMBOL_GO)[0]
else:
go_id = self.data.vocab_simple.encode(constant.SYMBOL_GO)
batch_go = tf.tile(
tf.expand_dims(self.embedding_fn(go_id, emb_simple), axis=0),
[self.model_config.batch_size, 1])
# For static_seq train_mode
if self.model_config.npad_mode == 'static_seq':
with tf.variable_scope('npad'):
npad_w = tf.get_variable(
'npad_w', shape=[1, self.model_config.dimension, self.model_config.dimension],
initializer=tf.contrib.layers.xavier_initializer(), dtype=tf.float32)
obj_tensors['npad_w'] = npad_w
if self.is_train and (train_mode == 'teacher' or
train_mode == 'teachercritical'or train_mode == 'teachercriticalv2'):
# General train
print('Use Generally Process.')
decoder_embed_inputs_list = self.embedding_fn(
sentence_simple_input_placeholder[:-1], emb_simple)
final_output, decoder_output, cur_context = self.decode_step(
decoder_embed_inputs_list, encoder_outputs, encoder_attn_bias,
rule_id_input_placeholder, mem_contexts, mem_outputs, global_step, score, batch_go,
obj_tensors)
decoder_logit = (
tf.nn.conv1d(final_output, tf.expand_dims(tf.transpose(w), axis=0), 1, 'SAME') +
tf.expand_dims(tf.expand_dims(b, axis=0), axis=0))
decoder_target_list = []
decoder_logit_list = tf.unstack(decoder_logit, axis=1)
for logit in decoder_logit_list:
decoder_target_list.append(tf.argmax(logit, output_type=tf.int32, axis=-1))
decoder_output_list = [
tf.squeeze(d, 1)
for d in tf.split(decoder_output, self.model_config.max_simple_sentence, axis=1)]
final_output_list = [
tf.squeeze(d, 1)
for d in tf.split(final_output, self.model_config.max_simple_sentence, axis=1)]
if self.model_config.pointer_mode:
segment_mask = None
if 'line_comp_segids' in obj:
segment_mask = obj['line_comp_segids']
decoder_logit_list = word_distribution(
decoder_logit_list, decoder_output_list, encoder_outputs, encoder_embed_inputs,
sentence_complex_input_placeholder, obj_tensors, self.model_config, self.data, segment_mask)
elif self.is_train and (train_mode == 'static_seq' or train_mode == 'static_self-critical'):
decoder_target_list = []
decoder_logit_list = []
decoder_embed_inputs_list = []
# Will Override for following 3 lists
final_output_list = []
decoder_output_list = []
contexts = []
sample_target_list = []
sample_logit_list = []
gpu_assign_interval = int(self.model_config.max_simple_sentence / 3)
for step in range(self.model_config.max_simple_sentence):
gpu_id = int(step / gpu_assign_interval)
if gpu_id > 3:
gpu_id = 3
gpu_id += 1
with tf.device('/device:GPU:%s' % gpu_id):
print('Step%s with GPU%s' % (step, gpu_id))
final_outputs, _, cur_context = self.decode_step(
decoder_embed_inputs_list, encoder_outputs, encoder_attn_bias,
rule_id_input_placeholder, mem_contexts, mem_outputs, global_step,
score, batch_go, obj_tensors)
final_output_list = [
tf.squeeze(d, 1)
for d in tf.split(final_outputs, step+1, axis=1)]
final_output = final_output_list[-1]
# if self.model_config.npad_mode == 'static_seq':
# final_output = tf.matmul(final_output, npad_w)
last_logit_list = self.output_to_logit(final_output, w, b)
last_target_list = tf.argmax(last_logit_list, output_type=tf.int32, axis=-1)
decoder_logit_list.append(last_logit_list)
decoder_target_list.append(last_target_list)
decoder_embed_inputs_list.append(
tf.stop_gradient(self.embedding_fn(last_target_list, emb_simple)))
if train_mode == 'static_self-critical':
last_sample_list = tf.multinomial(last_logit_list, 1)
sample_target_list.append(last_sample_list)
indices = tf.stack(
[tf.range(0, self.model_config.batch_size, dtype=tf.int64),
tf.squeeze(last_sample_list)],
axis=-1)
sample_logit_list.append(tf.gather_nd(tf.nn.softmax(last_logit_list), indices))
else:
# Beam Search
print('Use Beam Search with Beam Search Size %d.' % self.model_config.beam_search_size)
return self.transformer_beam_search(encoder_outputs, encoder_attn_bias, encoder_embed_inputs_list,
sentence_complex_input_placeholder, emb_simple, w, b,
rule_id_input_placeholder, mem_contexts, mem_outputs, global_step,
score, obj, obj_tensors)
gt_target_list = sentence_simple_input_placeholder
output = ModelOutput(
contexts=cur_context if 'rule' in self.model_config.memory else None,
encoder_outputs=encoder_outputs,
decoder_outputs_list=final_output_list if train_mode != 'dynamic_self-critical' else None,
final_outputs_list=final_output_list if train_mode != 'dynamic_self-critical' else None,
decoder_logit_list=decoder_logit_list if train_mode != 'dynamic_self-critical' else None,
gt_target_list=gt_target_list,
encoder_embed_inputs_list=tf.unstack(encoder_embed_inputs, axis=1),
decoder_target_list=decoder_target_list,
sample_logit_list=sampled_logit_list if train_mode == 'dynamic_self-critical' else None,
sample_target_list=sampled_target_list if train_mode == 'dynamic_self-critical' else None,
pred_score_tuple=pred_score_tuple if 'pred' in self.model_config.tune_mode else None,
obj_tensors=obj_tensors,
)
return output
def _prepare_decoder_input(area_encoding,
decoder_nonpadding,
features,
hparams,
embed_scope=None):
"""Prepare the input for the action decoding.
Args:
area_encoding: the encoder output in shape of [batch_size, area_len, depth].
decoder_nonpadding: the nonpadding mask for the decoding seq.
features: a dictionary of tensors in the shape of [batch_size, seq_length].
hparams: the hyperparameters.
embed_scope: the embedding scope.
Returns:
decoder_input: decoder input in shape of
[batch_size, num_steps, latent_depth]
decoder_self_attention_bias: decoder attention bias.
"""
with tf.variable_scope("prepare_decoder_input", reuse=tf.AUTO_REUSE):
shape = common_layers.shape_list(features["task"])
batch_size = shape[0]
encoder_input_length = shape[1]
depth = common_layers.shape_list(area_encoding)[-1]
if hparams.span_aggregation == "sum":
verb_embeds = span_embedding(encoder_input_length, area_encoding,
features["verb_refs"], hparams)
object_embeds = span_embedding(encoder_input_length, area_encoding,
features["obj_refs"], hparams)
input_embeds = span_embedding(encoder_input_length, area_encoding,
features["input_refs"], hparams)
non_input_embeds = tf.tile(
tf.expand_dims(
tf.expand_dims(
tf.get_variable(name="non_input_embeds",
shape=[depth]), 0), 0),
[batch_size,
tf.shape(features["input_refs"])[1], 1])
input_embeds = tf.where(
tf.tile(
tf.expand_dims(
tf.equal(features["input_refs"][:, :, 1],
features["input_refs"][:, :, 0]), 2),
[1, 1, tf.shape(input_embeds)[-1]]), non_input_embeds,
input_embeds)
elif hparams.span_aggregation == "mean":
area_encoding = area_encoding[:, :encoder_input_length, :]
verb_embeds = span_average_embed(area_encoding,
features["verb_refs"],
embed_scope)
object_embeds = span_average_embed(area_encoding,
features["obj_refs"],
embed_scope)
input_embeds = span_average_embed(area_encoding,
features["input_refs"],
embed_scope)
else:
raise ValueError("Unrecognized span aggregation method %s" %
(hparams.span_aggregation))
embeds = verb_embeds + object_embeds + input_embeds
embeds = tf.multiply(tf.expand_dims(decoder_nonpadding, 2), embeds)
start_embed = tf.tile(
tf.expand_dims(
tf.expand_dims(
tf.get_variable(name="start_step_embed", shape=[depth]),
0), 0), [batch_size, 1, 1])
embeds = tf.concat([start_embed, embeds], axis=1)
embeds = embeds[:, :-1, :]
decoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(
common_layers.shape_list(features["verb_refs"])[1]))
if hparams.pos == "timing":
decoder_input = common_attention.add_timing_signal_1d(embeds)
elif hparams.pos == "emb":
decoder_input = common_attention.add_positional_embedding(
embeds, hparams.max_length, "targets_positional_embedding",
None)
else:
decoder_input = embeds
return decoder_input, decoder_self_attention_bias
def _fast_decode(self,
features,
decode_length,
beam_size=1,
top_beams=1,
alpha=1.0):
"""Fast decoding.
Implements both greedy and beam search decoding, uses beam search iff
beam_size > 1, otherwise beam search related arguments are ignored.
Args:
features: a map of string to model features.
decode_length: an integer. How many additional timesteps to decode.
beam_size: number of beams.
top_beams: an integer. How many of the beams to return.
alpha: Float that controls the length penalty. larger the alpha, stronger
the preference for longer translations.
Returns:
A dict of decoding results {
"outputs": integer `Tensor` of decoded ids of shape
[batch_size, <= decode_length] if beam_size == 1 or
[batch_size, top_beams, <= decode_length]
"scores": decoding log probs from the beam search,
None if using greedy decoding (beam_size=1)
}
Raises:
NotImplementedError: If there are multiple data shards.
"""
if self._num_datashards != 1:
raise NotImplementedError(
"Fast decoding only supports a single shard.")
dp = self._data_parallelism
hparams = self._hparams
target_modality = self._problem_hparams.target_modality
if "targets_segmentation" in features:
raise NotImplementedError(
"Decoding not supported on packed datasets "
" If you want to decode from a dataset, use the non-packed version"
" of the dataset when decoding.")
if self.has_input:
inputs = features["inputs"]
if target_modality.is_class_modality:
decode_length = 1
else:
decode_length = (common_layers.shape_list(inputs)[1] +
features.get("decode_length", decode_length))
# TODO(llion): Clean up this reshaping logic.
inputs = tf.expand_dims(inputs, axis=1)
if len(inputs.shape) < 5:
inputs = tf.expand_dims(inputs, axis=4)
s = common_layers.shape_list(inputs)
batch_size = s[0]
inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
inputs = tf.squeeze(inputs, (2, 3))
# _shard_features called to ensure that the variable names match
inputs = self._shard_features({"inputs": inputs})["inputs"]
# input_modality = self._problem_hparams.input_modality["inputs"]
# with tf.variable_scope(input_modality.name):
# inputs = input_modality.bottom_sharded(inputs, dp)
encoder_output, encoder_decoder_attention_bias = dp(
self.encode, inputs, hparams, features=features)
encoder_output = encoder_output[0]
encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
partial_targets = None
else:
# The problem has no inputs.
encoder_output = None
encoder_decoder_attention_bias = None
# Prepare partial targets.
# In either features["inputs"] or features["targets"].
# We force the outputs to begin with these sequences.
partial_targets = features.get("inputs")
if partial_targets is None:
partial_targets = features["targets"]
assert partial_targets is not None
partial_targets = common_layers.expand_squeeze_to_nd(
partial_targets, 2)
partial_targets = tf.to_int64(partial_targets)
partial_targets_shape = common_layers.shape_list(partial_targets)
partial_targets_length = partial_targets_shape[1]
decode_length = (partial_targets_length +
features.get("decode_length", decode_length))
batch_size = partial_targets_shape[0]
with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
positional_encoding = common_attention.add_positional_embedding(
tf.zeros([1, decode_length, hparams.d_model]),
hparams.max_length, "positional_embedding", None)
def preprocess_targets(targets, i):
"""Performs preprocessing steps on the targets to prepare for the decoder.
This includes:
- Embedding the ids.
- Flattening to 3D tensor.
- Optionally adding timing signals.
Args:
targets: inputs ids to the decoder. [batch_size, 1]
i: scalar, Step number of the decoding loop.
Returns:
Processed targets [batch_size, 1, hidden_dim]
"""
targets_emb_var = self._get_targets_emb_var
targets = tf.gather(targets_emb_var, targets)
tf.logging.info("targets = %s" % targets)
targets = tf.squeeze(targets, (2, 3))
if positional_encoding is not None:
targets += positional_encoding[:, i:i + 1]
return targets
def symbols_to_logits_fn(ids, i, cache):
"""Go from ids to logits for next symbol."""
ids = ids[:, -1:]
targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
targets = preprocess_targets(targets, i)
bias = None # decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
body_outputs = dp(self.decode, targets,
cache.get("encoder_output"),
cache.get("encoder_decoder_attention_bias"),
bias, hparams, cache)
logits = body_outputs[0]
# with tf.variable_scope(target_modality.name):
# logits = target_modality.top_sharded(body_outputs, None, dp)[0]
ret = tf.squeeze(logits, axis=[1, 2, 3])
if partial_targets is not None:
# If the position is within the given partial targets, we alter the
# logits to always return those values.
# A faster approach would be to process the partial targets in one
# iteration in order to fill the corresponding parts of the cache.
# This would require broader changes, though.
vocab_size = tf.shape(ret)[1]
def forced_logits():
return tf.one_hot(
tf.tile(partial_targets[:, i], [beam_size]),
vocab_size, 0.0, -1e9)
ret = tf.cond(tf.less(i, partial_targets_length),
forced_logits, lambda: ret)
return ret, cache
ret = fast_decode(
encoder_output=encoder_output,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
symbols_to_logits_fn=symbols_to_logits_fn,
hparams=hparams,
decode_length=decode_length,
vocab_size=target_modality.top_dimensionality,
beam_size=beam_size,
top_beams=top_beams,
alpha=alpha,
batch_size=batch_size,
force_decode_length=self._decode_hparams.force_decode_length)
if partial_targets is not None:
if beam_size <= 1 or top_beams <= 1:
ret["outputs"] = ret["outputs"][:, partial_targets_length:]
else:
ret["outputs"] = ret["outputs"][:, :, partial_targets_length:]
return ret
def transformer_prepare_encoder(inputs,
target_space,
hparams,
features=None,
type_ids=None,
num_types=None,
reuse_target_embedding=tf.AUTO_REUSE):
"""Prepare one shard of the model for the encoder.
Args:
inputs: a Tensor.
target_space: a Tensor.
hparams: run hyperparameters
features: optionally pass the entire features dictionary as well.
This is needed now for "packed" datasets.
type_ids: optional, an int64 Tensor of shape [batch, length] that allows
for adding type embeddings, similar to positional embeddings.
num_types: optional, an int that decides the number of types in type_ids.
reuse_target_embedding: option to reuse variable name in the case that
symbol modalities are reused between inputs/targets.
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
attention
"""
ishape_static = inputs.shape.as_list()
encoder_input = inputs
if features and "inputs_segmentation" in features:
# Packed dataset. Keep the examples from seeing each other.
inputs_segmentation = features["inputs_segmentation"]
inputs_position = features["inputs_position"]
targets_segmentation = features["targets_segmentation"]
if (hasattr(hparams, "unidirectional_encoder")
and hparams.unidirectional_encoder):
tf.logging.info("Using unidirectional encoder")
encoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(
common_layers.shape_list(inputs)[1]))
else:
encoder_self_attention_bias = (
common_attention.attention_bias_same_segment(
inputs_segmentation, inputs_segmentation))
encoder_decoder_attention_bias = (
common_attention.attention_bias_same_segment(
targets_segmentation, inputs_segmentation))
else:
encoder_padding = common_attention.embedding_to_padding(encoder_input)
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
if (hasattr(hparams, "unidirectional_encoder")
and hparams.unidirectional_encoder):
tf.logging.info("Using unidirectional encoder")
encoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(
common_layers.shape_list(inputs)[1]))
else:
# Usual case - not a packed dataset.
encoder_self_attention_bias = ignore_padding
encoder_decoder_attention_bias = ignore_padding
inputs_position = None
if hparams.proximity_bias:
encoder_self_attention_bias += common_attention.attention_bias_proximal(
common_layers.shape_list(inputs)[1])
if target_space is not None and hparams.get("use_target_space_embedding",
True):
# Append target_space_id embedding to inputs.
emb_target_space = common_layers.embedding(
target_space,
32,
ishape_static[-1],
name="target_space_embedding",
dtype=hparams.get("activation_dtype", "float32"),
reuse=reuse_target_embedding)
emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
encoder_input += emb_target_space
if hparams.pos == "timing":
if inputs_position is not None:
encoder_input = common_attention.add_timing_signal_1d_given_position(
encoder_input, inputs_position)
else:
encoder_input = common_attention.add_timing_signal_1d(
encoder_input)
elif hparams.pos == "timing_from_features":
encoder_input = common_attention.add_timing_signals_from_features(
encoder_input, features, hparams.position_features)
elif hparams.pos == "emb":
encoder_input = common_attention.add_positional_embedding(
encoder_input, hparams.max_length, "inputs_positional_embedding",
inputs_position)
# Add type embeddings
if type_ids is not None:
if not num_types:
raise ValueError("Need to set num_types as well.")
encoder_input = common_attention.add_positional_embedding(
encoder_input, num_types, "inputs_type_embedding", type_ids)
encoder_self_attention_bias = common_layers.cast_like(
encoder_self_attention_bias, encoder_input)
encoder_decoder_attention_bias = common_layers.cast_like(
encoder_decoder_attention_bias, encoder_input)
return (encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias)
def _fast_decode(self,
features,
decode_length,
beam_size=1,
top_beams=1,
alpha=1.0):
"""Fast decoding.
Overrides tensor2tensor.models.transformer.Transformer._fast_decode
to let symbols_to_logits_fn return multiple things.
Implements both greedy and beam search decoding, uses beam search iff
beam_size > 1, otherwise beam search related arguments are ignored.
Args:
features: a map of string to model features.
decode_length: an integer. How many additional timesteps to decode.
beam_size: number of beams.
top_beams: an integer. How many of the beams to return.
alpha: Float that controls the length penalty. larger the alpha, stronger
the preference for longer translations.
Returns:
A dict of decoding results {
"body_output": tensor of size
[batch_size, <= decode_length, hidden_size]
(or [batch_size, top_beams, <= decode_length, hidden_size])
giving the raw output of the Transformer decoder corresponding
to the predicted sequences
"outputs": integer `Tensor` of decoded ids of shape
[batch_size, <= decode_length] if beam_size == 1 or
[batch_size, top_beams, <= decode_length]
"scores": decoding log probs from the beam search,
None if using greedy decoding (beam_size=1)
}
Raises:
NotImplementedError: If there are multiple data shards.
"""
if self._num_datashards != 1:
raise NotImplementedError(
"Fast decoding only supports a single shard.")
dp = self._data_parallelism
hparams = self._hparams
target_modality = self._problem_hparams.target_modality
if isinstance(target_modality, dict):
primary_target_feature = self._problem_hparams.primary_target_modality
primary_target_modality = target_modality[primary_target_feature]
bottom_variable_scope = "%s/%s" % (primary_target_modality.name,
primary_target_feature)
else:
primary_target_feature = "targets"
primary_target_modality = target_modality
bottom_variable_scope = target_modality.name
if self.has_input:
inputs = features["inputs"]
if primary_target_modality.is_class_modality:
decode_length = 1
else:
decode_length = (common_layers.shape_list(inputs)[1] +
features.get("decode_length", decode_length))
# TODO(llion): Clean up this reshaping logic.
inputs = tf.expand_dims(inputs, axis=1)
if len(inputs.shape) < 5:
inputs = tf.expand_dims(inputs, axis=4)
s = common_layers.shape_list(inputs)
batch_size = s[0]
inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
# _shard_features called to ensure that the variable names match
inputs = self._shard_features({"inputs": inputs})["inputs"]
input_modality = self._problem_hparams.input_modality["inputs"]
with tf.variable_scope(input_modality.name):
inputs = input_modality.bottom_sharded(inputs, dp)
with tf.variable_scope("body"):
encoder_output, encoder_decoder_attention_bias = dp(
self.encode,
inputs,
features["target_space_id"],
hparams,
features=features)
encoder_output = encoder_output[0]
encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
partial_targets = None
else:
# The problem has no inputs.
encoder_output = None
encoder_decoder_attention_bias = None
# Prepare partial targets.
# In either features["inputs"] or features["targets"].
# We force the outputs to begin with these sequences.
partial_targets = features.get("inputs")
if partial_targets is None:
partial_targets = features[primary_target_feature]
assert partial_targets is not None
partial_targets = common_layers.expand_squeeze_to_nd(
partial_targets, 2)
partial_targets = tf.to_int64(partial_targets)
partial_targets_shape = common_layers.shape_list(partial_targets)
partial_targets_length = partial_targets_shape[1]
decode_length = (partial_targets_length +
features.get("decode_length", decode_length))
batch_size = partial_targets_shape[0]
if hparams.pos == "timing":
positional_encoding = common_attention.get_timing_signal_1d(
decode_length + 1, hparams.hidden_size)
elif hparams.pos == "emb":
positional_encoding = common_attention.add_positional_embedding(
tf.zeros([1, decode_length + 1, hparams.hidden_size]),
hparams.max_length, "targets_positional_embedding", None)
else:
positional_encoding = None
def preprocess_targets(targets, i):
"""Performs preprocessing steps on the targets to prepare for the decoder.
This includes:
- Embedding the ids.
- Flattening to 3D tensor.
- Optionally adding timing signals.
Args:
targets: inputs ids to the decoder. [batch_size, 1]
i: scalar, Step number of the decoding loop.
Returns:
Processed targets [batch_size, 1, hidden_dim]
"""
# _shard_features called to ensure that the variable names match
targets = self._shard_features({primary_target_feature:
targets})[primary_target_feature]
with tf.variable_scope(bottom_variable_scope):
targets = primary_target_modality.targets_bottom_sharded(
targets, dp)[0]
targets = common_layers.flatten4d3d(targets)
# At step 0, targets will have 0 size, and instead we want to
# create an embedding of all-zero, corresponding to the start symbol
# this matches what transformer_prepare_decoder does to the target
# outputs during training
targets = tf.cond(tf.equal(i, 0), lambda: tf.zeros_like(targets),
lambda: targets)
if positional_encoding is not None:
targets += positional_encoding[:, i:i + 1]
return targets
decoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(decode_length))
if hparams.proximity_bias:
decoder_self_attention_bias += common_attention.attention_bias_proximal(
decode_length)
def symbols_to_logits_fn(ids, i, cache):
"""Go from ids to logits for next symbol."""
ids = ids[:, -1:]
targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
targets = preprocess_targets(targets, i)
bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
logits = self._symbols_to_logits_fn(targets, features, bias, cache)
logits = tf.squeeze(logits, axis=[1, 2, 3])
if partial_targets is not None:
# If the position is within the given partial targets, we alter the
# logits to always return those values.
# A faster approach would be to process the partial targets in one
# iteration in order to fill the corresponding parts of the cache.
# This would require broader changes, though.
vocab_size = tf.shape(logits)[1]
def forced_logits():
return tf.one_hot(
tf.tile(partial_targets[:, i], [beam_size]),
vocab_size, 0.0, -1e9)
logits = tf.cond(tf.less(i, partial_targets_length),
forced_logits, lambda: logits)
return logits, cache
cache = dict()
infer_out = dict()
if encoder_output is not None:
padding_mask = 1. - common_attention.attention_bias_to_padding(
encoder_decoder_attention_bias)
masked_encoded_output = encoder_output * tf.expand_dims(
padding_mask, axis=2)
infer_out["encoded_inputs"] = tf.reduce_sum(masked_encoded_output,
axis=1)
self._prepare_decoder_cache(batch_size, features, cache)
ret = fast_decode(
encoder_output=encoder_output,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
symbols_to_logits_fn=symbols_to_logits_fn,
hparams=hparams,
decode_length=decode_length,
vocab_size=primary_target_modality.top_dimensionality,
beam_size=beam_size,
top_beams=top_beams,
alpha=alpha,
batch_size=batch_size,
force_decode_length=self._decode_hparams.force_decode_length,
cache=cache)
infer_out.update(ret)
if "cache" in ret:
infer_out.update(ret["cache"])
if partial_targets is not None:
if beam_size <= 1 or top_beams <= 1:
infer_out["outputs"] = infer_out[
"outputs"][:, partial_targets_length:]
else:
infer_out["outputs"] = infer_out[
"outputs"][:, :, partial_targets_length:]
return infer_out
def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
"""Prepare one shard of the model for the encoder.
Args:
inputs: a Tensor.
target_space: a Tensor.
hparams: run hyperparameters
features: optionally pass the entire features dictionary as well.
This is needed now for "packed" datasets.
Returns:
encoder_input: a Tensor, bottom of encoder stack
encoder_self_attention_bias: a bias tensor for use in encoder self-attention
encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
attention
"""
ishape_static = inputs.shape.as_list()
encoder_input = inputs
if features and "inputs_segmentation" in features:
# Packed dataset. Keep the examples from seeing each other.
inputs_segmentation = features["inputs_segmentation"]
inputs_position = features["inputs_position"]
targets_segmentation = features["targets_segmentation"]
encoder_self_attention_bias = common_attention.attention_bias_same_segment(
inputs_segmentation, inputs_segmentation)
encoder_decoder_attention_bias = (
common_attention.attention_bias_same_segment(targets_segmentation,
inputs_segmentation))
else:
# Usual case - not a packed dataset.
encoder_padding = common_attention.embedding_to_padding(encoder_input)
ignore_padding = common_attention.attention_bias_ignore_padding(
encoder_padding)
encoder_self_attention_bias = ignore_padding
encoder_decoder_attention_bias = ignore_padding
inputs_position = None
if hparams.proximity_bias:
encoder_self_attention_bias += common_attention.attention_bias_proximal(
common_layers.shape_list(inputs)[1])
if hparams.get("use_target_space_embedding", True):
# Append target_space_id embedding to inputs.
emb_target_space = common_layers.embedding(
target_space,
32,
ishape_static[-1],
name="target_space_embedding",
dtype=tf.bfloat16
if hparams.activation_dtype == "bfloat16" else tf.float32)
emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
encoder_input += emb_target_space
if hparams.pos == "timing":
if inputs_position is not None:
encoder_input = common_attention.add_timing_signal_1d_given_position(
encoder_input, inputs_position)
else:
encoder_input = common_attention.add_timing_signal_1d(encoder_input)
elif hparams.pos == "emb":
encoder_input = common_attention.add_positional_embedding(
encoder_input, hparams.max_length, "inputs_positional_embedding",
inputs_position)
if hparams.activation_dtype == "bfloat16":
encoder_self_attention_bias = tf.cast(encoder_self_attention_bias,
tf.bfloat16)
encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
tf.bfloat16)
return (encoder_input, encoder_self_attention_bias,
encoder_decoder_attention_bias)
def _fast_decode(self,
features,
decode_length,
beam_size=1,
top_beams=1,
alpha=1.0):
"""Fast decoding.
Implements both greedy and beam search decoding, uses beam search iff
beam_size > 1, otherwise beam search related arguments are ignored.
Args:
features: a map of string to model features.
decode_length: an integer. How many additional timesteps to decode.
beam_size: number of beams.
top_beams: an integer. How many of the beams to return.
alpha: Float that controls the length penalty. larger the alpha, stronger
the preference for longer translations.
Returns:
A dict of decoding results {
"outputs": integer `Tensor` of decoded ids of shape
[batch_size, <= decode_length] if beam_size == 1 or
[batch_size, top_beams, <= decode_length]
"scores": decoding log probs from the beam search,
None if using greedy decoding (beam_size=1)
}
Raises:
NotImplementedError: If there are multiple data shards.
"""
if self._num_datashards != 1:
raise NotImplementedError("Fast decoding only supports a single shard.")
dp = self._data_parallelism
hparams = self._hparams
target_modality = self._problem_hparams.modality["targets"]
target_vocab_size = self._problem_hparams.vocab_size["targets"]
if target_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
target_vocab_size += (-target_vocab_size) % hparams.vocab_divisor
if "targets_segmentation" in features:
raise NotImplementedError(
"Decoding not supported on packed datasets "
" If you want to decode from a dataset, use the non-packed version"
" of the dataset when decoding.")
if self.has_input:
inputs = features["inputs"]
if target_modality == modalities.ModalityType.CLASS_LABEL:
decode_length = 1
else:
decode_length = (
common_layers.shape_list(inputs)[1] + features.get(
"decode_length", decode_length))
# TODO(llion): Clean up this reshaping logic.
inputs = tf.expand_dims(inputs, axis=1)
if len(inputs.shape) < 5:
inputs = tf.expand_dims(inputs, axis=4)
s = common_layers.shape_list(inputs)
batch_size = s[0]
inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
# _shard_features called to ensure that the variable names match
inputs = self._shard_features({"inputs": inputs})["inputs"]
input_modality = self._problem_hparams.modality["inputs"]
input_vocab_size = self._problem_hparams.vocab_size["inputs"]
if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
modality_name = hparams.name.get(
"inputs",
modalities.get_name(input_modality))(hparams, input_vocab_size)
with tf.variable_scope(modality_name):
bottom = hparams.bottom.get("inputs",
modalities.get_bottom(input_modality))
inputs = dp(bottom, inputs, hparams, input_vocab_size)
with tf.variable_scope("body"):
encoder_output, encoder_decoder_attention_bias = dp(
self.encode,
inputs,
features["target_space_id"],
hparams,
features=features)
encoder_output = encoder_output[0]
encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
if 'partial_targets' in features:
partial_targets = features['partial_targets']
else:
partial_targets = None
else:
# The problem has no inputs.
encoder_output = None
encoder_decoder_attention_bias = None
# Prepare partial targets.
# In either features["inputs"] or features["targets"].
# We force the outputs to begin with these sequences.
partial_targets = features.get("inputs")
if partial_targets is None:
partial_targets = features["targets"]
assert partial_targets is not None
if partial_targets is not None:
partial_targets = common_layers.expand_squeeze_to_nd(partial_targets, 2)
partial_targets = tf.to_int64(partial_targets)
partial_targets_shape = common_layers.shape_list(partial_targets)
partial_targets_length = partial_targets_shape[1]
decode_length = (
partial_targets_length + features.get("decode_length", decode_length))
batch_size = partial_targets_shape[0]
if hparams.pos == "timing":
positional_encoding = common_attention.get_timing_signal_1d(
decode_length + 1, hparams.hidden_size)
elif hparams.pos == "emb":
positional_encoding = common_attention.add_positional_embedding(
tf.zeros([1, decode_length, hparams.hidden_size]), hparams.max_length,
"body/targets_positional_embedding", None)
else:
positional_encoding = None
def preprocess_targets(targets, i):
"""Performs preprocessing steps on the targets to prepare for the decoder.
This includes:
- Embedding the ids.
- Flattening to 3D tensor.
- Optionally adding timing signals.
Args:
targets: inputs ids to the decoder. [batch_size, 1]
i: scalar, Step number of the decoding loop.
Returns:
Processed targets [batch_size, 1, hidden_dim]
"""
# _shard_features called to ensure that the variable names match
targets = self._shard_features({"targets": targets})["targets"]
modality_name = hparams.name.get(
"targets",
modalities.get_name(target_modality))(hparams, target_vocab_size)
with tf.variable_scope(modality_name):
bottom = hparams.bottom.get(
"targets", modalities.get_targets_bottom(target_modality))
targets = dp(bottom, targets, hparams, target_vocab_size)[0]
targets = common_layers.flatten4d3d(targets)
# GO embeddings are all zero, this is because transformer_prepare_decoder
# Shifts the targets along by one for the input which pads with zeros.
# If the modality already maps GO to the zero embeddings this is not
# needed.
targets = tf.cond(
tf.equal(i, 0), lambda: tf.zeros_like(targets), lambda: targets)
if positional_encoding is not None:
targets += positional_encoding[:, i:i + 1]
return targets
decoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(decode_length))
if hparams.proximity_bias:
decoder_self_attention_bias += common_attention.attention_bias_proximal(
decode_length)
# Create tensors for encoder-decoder attention history
att_cache = {"attention_history": {}}
num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
att_batch_size, enc_seq_length = common_layers.shape_list(encoder_output)[0:2]
for layer in range(num_layers):
att_cache["attention_history"]["layer_%d" % layer] = tf.zeros(
[att_batch_size, hparams.num_heads, 0, enc_seq_length])
att_cache["body_outputs"] = tf.zeros([att_batch_size, 1, 0, hparams.hidden_size])
def update_decoder_attention_history(cache):
for k in filter(lambda x: "decoder" in x and not "self" in x and not "logits" in x,
self.attention_weights.keys()):
m = re.search(r"(layer_\d+)", k)
if m is None:
continue
cache["attention_history"][m[0]] = tf.concat(
[cache["attention_history"][m[0]], self.attention_weights[k]], axis=2)
def symbols_to_logits_fn(ids, i, cache):
"""Go from ids to logits for next symbol."""
ids = ids[:, -1:]
targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
targets = preprocess_targets(targets, i)
bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]
with tf.variable_scope("body"):
body_outputs = dp(
self.decode,
targets,
cache.get("encoder_output"),
cache.get("encoder_decoder_attention_bias"),
bias,
hparams,
cache,
nonpadding=features_to_nonpadding(features, "targets"))
update_decoder_attention_history(cache)
cache["body_outputs"] = tf.concat([cache["body_outputs"], body_outputs[0]], axis=2)
modality_name = hparams.name.get(
"targets",
modalities.get_name(target_modality))(hparams, target_vocab_size)
with tf.variable_scope(modality_name):
top = hparams.top.get("targets", modalities.get_top(target_modality))
logits = dp(top, body_outputs, None, hparams, target_vocab_size)[0]
ret = tf.squeeze(logits, axis=[1, 2, 3])
if partial_targets is not None:
# If the position is within the given partial targets, we alter the
# logits to always return those values.
# A faster approach would be to process the partial targets in one
# iteration in order to fill the corresponding parts of the cache.
# This would require broader changes, though.
vocab_size = tf.shape(ret)[1]
def forced_logits():
return tf.one_hot(
tf.tile(partial_targets[:, i], [beam_size]), vocab_size, 0.0,
-1e9)
ret = tf.cond(
tf.less(i, partial_targets_length), forced_logits, lambda: ret)
return ret, cache
ret = fast_decode(
encoder_output=encoder_output,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
symbols_to_logits_fn=symbols_to_logits_fn,
hparams=hparams,
decode_length=decode_length,
vocab_size=target_vocab_size,
beam_size=beam_size,
top_beams=top_beams,
alpha=alpha,
batch_size=batch_size,
force_decode_length=self._decode_hparams.force_decode_length,
cache=att_cache)
if partial_targets is not None:
if beam_size <= 1 or top_beams <= 1:
ret["outputs"] = ret["outputs"][:, partial_targets_length:]
else:
ret["outputs"] = ret["outputs"][:, :, partial_targets_length:]
return ret
def _fast_decode_tpu(self, features, decode_length, beam_size=1):
"""Fast decoding.
Implements only greedy decoding on TPU.
Args:
features: A map of string to model features.
decode_length: An integer, how many additional timesteps to decode.
beam_size: An integer, number of beams.
Returns:
A dict of decoding results {
"outputs": integer `Tensor` of decoded ids of shape
[batch_size, <= decode_length]
"scores": decoding log probs from the beam search,
None if using greedy decoding (beam_size=1)
}.
Raises:
NotImplementedError: If there are multiple data shards or beam_size > 1.
"""
if self._num_datashards != 1:
raise NotImplementedError(
"Fast decoding only supports a single shard.")
if "targets_segmentation" in features:
raise NotImplementedError(
"Decoding not supported on packed datasets "
" If you want to decode from a dataset, use the non-packed version"
" of the dataset when decoding.")
dp = self._data_parallelism
hparams = self._hparams
target_modality = self._problem_hparams.target_modality
if self.has_input:
inputs = features["inputs"]
if target_modality.is_class_modality:
decode_length = 1
else:
decode_length = (common_layers.shape_list(inputs)[1] +
features.get("decode_length", decode_length))
# TODO(llion): Clean up this reshaping logic.
inputs = tf.expand_dims(inputs, axis=1)
if len(inputs.shape) < 5:
inputs = tf.expand_dims(inputs, axis=4)
s = common_layers.shape_list(inputs)
batch_size = s[0]
inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
# _shard_features called to ensure that the variable names match
inputs = self._shard_features({"inputs": inputs})["inputs"]
input_modality = self._problem_hparams.input_modality["inputs"]
with tf.variable_scope(input_modality.name):
inputs = input_modality.bottom_sharded(inputs, dp)
with tf.variable_scope("body"):
encoder_output, encoder_decoder_attention_bias = dp(
self.encode,
inputs,
features["target_space_id"],
hparams,
features=features)
encoder_output = encoder_output[0]
encoder_decoder_attention_bias = encoder_decoder_attention_bias[0]
partial_targets = None
else:
# The problem has no inputs.
encoder_output = None
encoder_decoder_attention_bias = None
# Prepare partial targets.
# In either features["inputs"] or features["targets"].
# We force the outputs to begin with these sequences.
partial_targets = features.get("inputs")
if partial_targets is None:
partial_targets = features["targets"]
assert partial_targets is not None
partial_targets = common_layers.expand_squeeze_to_nd(
partial_targets, 2)
partial_targets = tf.to_int64(partial_targets)
partial_targets_shape = common_layers.shape_list(partial_targets)
partial_targets_length = partial_targets_shape[1]
decode_length = (partial_targets_length +
features.get("decode_length", decode_length))
batch_size = partial_targets_shape[0]
if hparams.pos == "timing":
positional_encoding = common_attention.get_timing_signal_1d(
decode_length + 1, hparams.hidden_size)
elif hparams.pos == "emb":
positional_encoding = common_attention.add_positional_embedding(
tf.zeros([1, decode_length + 1, hparams.hidden_size]),
hparams.max_length, "body/targets_positional_embedding", None)
else:
positional_encoding = None
def preprocess_targets(targets, i):
"""Performs preprocessing steps on the targets to prepare for the decoder.
This includes:
- Embedding the ids.
- Flattening to 3D tensor.
- Optionally adding timing signals.
Args:
targets: A tensor, inputs ids to the decoder. [batch_size, 1].
i: An integer, Step number of the decoding loop.
Returns:
A tensor, processed targets [batch_size, 1, hidden_dim].
"""
# _shard_features called to ensure that the variable names match
targets = self._shard_features({"targets": targets})["targets"]
with tf.variable_scope(target_modality.name):
targets = target_modality.targets_bottom_sharded(targets,
dp)[0]
targets = common_layers.flatten4d3d(targets)
# TODO(llion): Explain! Is this even needed?
targets = tf.cond(tf.equal(i, 0), lambda: tf.zeros_like(targets),
lambda: targets)
if positional_encoding is not None:
positional_encoding_shape = positional_encoding.shape.as_list()
targets += tf.slice(positional_encoding, [0, i, 0], [
positional_encoding_shape[0], 1,
positional_encoding_shape[2]
])
return targets
decoder_self_attention_bias = (
common_attention.attention_bias_lower_triangle(decode_length))
if hparams.proximity_bias:
decoder_self_attention_bias += common_attention.attention_bias_proximal(
decode_length)
def symbols_to_logits_tpu_fn(ids, i, cache):
"""Go from ids to logits for next symbol on TPU.
Args:
ids: A tensor, symbol IDs.
i: An integer, step number of the decoding loop. Only used for inference
on TPU.
cache: A dict, containing tensors which are the results of previous
attentions, used for fast decoding.
Returns:
ret: A tensor, computed logits.
cache: A dict, containing tensors which are the results of previous
attentions, used for fast decoding.
"""
ids = ids[:, -1:]
targets = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
targets = preprocess_targets(targets, i)
bias_shape = decoder_self_attention_bias.shape.as_list()
bias = tf.slice(decoder_self_attention_bias, [0, 0, i, 0],
[bias_shape[0], bias_shape[1], 1, bias_shape[3]])
with tf.variable_scope("body"):
body_outputs = dp(self.decode,
targets,
cache.get("encoder_output"),
cache.get("encoder_decoder_attention_bias"),
bias,
hparams,
cache,
i,
nonpadding=features_to_nonpadding(
features, "targets"))
with tf.variable_scope(target_modality.name):
logits = target_modality.top_sharded(body_outputs, None, dp)[0]
ret = tf.squeeze(logits, axis=[1, 2, 3])
if partial_targets is not None:
# If the position is within the given partial targets, we alter the
# logits to always return those values.
# A faster approach would be to process the partial targets in one
# iteration in order to fill the corresponding parts of the cache.
# This would require broader changes, though.
vocab_size = tf.shape(ret)[1]
def forced_logits():
return tf.one_hot(
tf.tile(
tf.slice(partial_targets, [0, i],
[partial_targets.shape.as_list()[0], 1]),
[beam_size]), vocab_size, 0.0, -1e9)
ret = tf.cond(tf.less(i, partial_targets_length),
forced_logits, lambda: ret)
return ret, cache
ret = fast_decode_tpu(
encoder_output=encoder_output,
encoder_decoder_attention_bias=encoder_decoder_attention_bias,
symbols_to_logits_fn=symbols_to_logits_tpu_fn,
hparams=hparams,
decode_length=decode_length,
beam_size=beam_size,
batch_size=batch_size,
force_decode_length=self._decode_hparams.force_decode_length)
if partial_targets is not None:
ret["outputs"] = ret["outputs"][:, partial_targets_length:]
return ret
