def padded_sequence_accuracy(predictions,
labels,
weights_fn=common_layers.weights_nonzero):
"""Percentage of times that predictions matches labels everywhere (non-0)."""
# If the last dimension is 1 then we're using L1/L2 loss.
if common_layers.shape_list(predictions)[-1] == 1:
return rounding_sequence_accuracy(
predictions, labels, weights_fn=weights_fn)
with tf.variable_scope(
"padded_sequence_accuracy", values=[predictions, labels]):
padded_predictions, padded_labels = common_layers.pad_with_zeros(
predictions, labels)
weights = weights_fn(padded_labels)
# Flatten, keeping batch dim (and num_classes dim for predictions)
# TPU argmax can only deal with a limited number of dimensions
predictions_shape = common_layers.shape_list(padded_predictions)
batch_size = predictions_shape[0]
num_classes = predictions_shape[-1]
flat_size = common_layers.list_product(
common_layers.shape_list(padded_labels)[1:])
padded_predictions = tf.reshape(
padded_predictions,
[batch_size, common_layers.list_product(predictions_shape[1:-1]),
num_classes])
padded_labels = tf.reshape(padded_labels, [batch_size, flat_size])
weights = tf.reshape(weights, [batch_size, flat_size])
outputs = tf.to_int32(tf.argmax(padded_predictions, axis=-1))
padded_labels = tf.to_int32(padded_labels)
not_correct = to_float(tf.not_equal(outputs, padded_labels)) * weights
axis = list(range(1, len(outputs.get_shape())))
correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
return correct_seq, tf.constant(1.0)
def reduce_dimensions(predictions, labels):
"""Reduce dimensions for high-dimensional predictions and labels."""
# We will treat first dimensions as batch. One example are video frames.
if len(predictions.get_shape()) > 5:
predictions_shape = common_layers.shape_list(predictions)
predictions = tf.reshape(
predictions, [predictions_shape[0], predictions_shape[1], -1,
predictions_shape[-1]])
labels_shape = common_layers.shape_list(labels)
labels = tf.reshape(
labels, [labels_shape[0], labels_shape[1], -1])
return predictions, labels
def body_fn(i, ss_tokens):
"""Constructs conditioning tokens for scheduled sampling."""
# next_token_logits depends on timesteps 0...i-1.
#
# [batch_size, seq_len] -> [batch_size, seq_len, vocab_size]
ss_tokens_logits = infer_fn(ss_tokens)
# Same as 'next_token_logits = ss_tokens_logits[:, i, :]'.
vocab_size = common_layers.shape_list(ss_tokens_logits)[2]
next_token_logits = tf.slice(ss_tokens_logits,
begin=[0, i, 0],
size=[batch_size, 1, vocab_size])
next_token_logits = tf.squeeze(next_token_logits, axis=[1])
# [batch_size, vocab_size] -> [batch_size]
sampled_next_tokens = _sample_next_tokens(next_token_logits)
# Same as 'gold_next_tokens = targets[:, i]'.
gold_next_tokens = tf.slice(targets,
begin=[0, i],
size=[batch_size, 1])
gold_next_tokens = tf.squeeze(gold_next_tokens, axis=[1])
next_tokens = mix_fn(gold_next_tokens, sampled_next_tokens)
ss_tokens = _update_timestep(ss_tokens, timestep=i, values=next_tokens)
return i + 1, tf.stop_gradient(ss_tokens)
def _sample_next_tokens(logits):
"""Sample tokens for next timestep."""
batch_size = common_layers.shape_list(logits)[0]
next_tokens = tf.random.categorical(logits, 1)
next_tokens = tf.cast(next_tokens, tf.int32)
next_tokens = tf.reshape(next_tokens, [batch_size])
return next_tokens
def infer_fn(self, partial_targets):
"""Computes logits for all timesteps.
Args:
partial_targets: [batch_size, seq_len]. Targets to condition on.
Returns:
next_token_logits: [batch_size, seq_len, vocab_size]
"""
batch_size, seq_len = common_layers.shape_list(partial_targets)
partial_targets = tf.reshape(partial_targets,
[batch_size, seq_len, 1, 1])
features = copy.copy(self._features)
features["targets"] = partial_targets
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
transformed_features = self._t2tmodel.bottom(features)
with tf.variable_scope("body"):
body_outputs, losses = self._t2tmodel._normalize_body_output( # pylint: disable=protected-access
self._t2tmodel.body(transformed_features))
assert losses == {
"extra": 0.0
}, ("Auxiliary losses are not propagated in this code. %s" %
(losses, ))
logits = self._t2tmodel.top(body_outputs, features)
vocab_size = self._t2tmodel.problem_hparams.vocab_size["targets"]
logits = tf.reshape(logits, [batch_size, seq_len, vocab_size])
return logits
def sequential_scheduled_sampling_for_t2tmodel(t2tmodel, features):
"""Schedule Sampling for T2TModels.
Args:
t2tmodel: T2TModel instance.
features: {str: Tensor}. Input features.
Returns:
ss_logits: [batch_size, seq_len, 1, 1, vocab_size].
losses_dict: {str: scalar Tensor}. Losses to minimize.
"""
targets = features["targets"]
targets_size = common_layers.shape_list(targets)
batch_size = targets_size[0]
seq_len = targets_size[1]
targets = tf.reshape(targets, [batch_size, seq_len])
adapter = ScheduledSamplingAdapter(t2tmodel, features)
ss_tokens, ss_logits, losses_dict = sequential_scheduled_sampling(
infer_fn=adapter.infer_fn,
mix_fn=adapter.mix_fn,
loss_fn=adapter.loss_fn,
targets=targets)
_ = ss_tokens # unused.
targets_vocab_size = t2tmodel.problem_hparams.vocab_size["targets"]
ss_logits = tf.reshape(ss_logits,
[batch_size, seq_len, 1, 1, targets_vocab_size])
return ss_logits, losses_dict
def image_rmse(predictions, labels, weights_fn=common_layers.weights_all):
"""RMSE but will argmax if last dim is not 1."""
if common_layers.shape_list(predictions)[-1] == 1:
predictions = tf.squeeze(predictions, axis=[-1])
else:
predictions = tf.argmax(predictions, axis=-1)
return padded_rmse(predictions, labels, weights_fn)
def symbol_top(body_output, targets, model_hparams, vocab_size):
"""Generate logits.
Args:
body_output: A Tensor with shape
[batch, p0, p1, model_hparams.hidden_size].
targets: Unused.
model_hparams: HParams, model hyperparmeters.
vocab_size: int, vocabulary size.
Returns:
logits: A Tensor with shape [batch, p0, p1, ?, vocab_size].
"""
del targets # unused arg
if model_hparams.shared_embedding_and_softmax_weights:
scope_name = "shared"
reuse = tf.AUTO_REUSE
else:
scope_name = "softmax"
reuse = False
with tf.variable_scope(scope_name, reuse=reuse):
body_output_shape = common_layers.shape_list(body_output)
var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
if (model_hparams.factored_logits
and model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
# insert channels dimension
body_output = tf.expand_dims(body_output, 3)
return common_layers.FactoredTensor(body_output, var)
else:
body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
logits = tf.matmul(body_output, var, transpose_b=True)
return tf.reshape(logits, body_output_shape[:-1] + [1, vocab_size])
def basic_pool(features,
max_area_width,
max_area_height=1,
height=1,
fn=tf.reduce_max,
name=None):
"""Pools for each area based on a given pooling function (fn).
Args:
features: a Tensor in a shape of [batch_size, height * width, depth].
max_area_width: the max width allowed for an area.
max_area_height: the max height allowed for an area.
height: the height of the image.
fn: the TF function for the pooling.
name: the namescope.
Returns:
pool_results: A Tensor of shape [batch_size, num_areas, depth]
area_heights: A Tensor of shape [batch_size, num_areas, 1]
area_widths: A Tensor of shape [batch_size, num_areas, 1]
"""
with tf.name_scope(name, default_name="basic_pool"):
feature_shape = common_layers.shape_list(features)
batch_size = feature_shape[0]
length = feature_shape[-2]
depth = feature_shape[-1]
width = length // height
features_2d = tf.reshape(features, [batch_size, height, width, depth])
height_list = []
width_list = []
pool_list = []
size_tensor = tf.ones_like(features_2d[:, :, :, 0], dtype=tf.int32)
for area_height in range(max_area_height):
for area_width in range(max_area_width):
pool_tensor = _pool_one_shape(features_2d,
area_width=area_width + 1,
area_height=area_height + 1,
batch_size=batch_size,
width=width,
height=height,
depth=depth,
fn=fn)
pool_list.append(
tf.reshape(pool_tensor, [batch_size, -1, depth]))
height_list.append(
tf.reshape(
size_tensor[:, area_height:, area_width:] *\
(area_height + 1), [batch_size, -1]))
width_list.append(
tf.reshape(
size_tensor[:, area_height:, area_width:] *\
(area_width + 1), [batch_size, -1]))
pool_results = tf.concat(pool_list, axis=1)
area_heights = tf.expand_dims(tf.concat(height_list, axis=1), 2)
area_widths = tf.expand_dims(tf.concat(width_list, axis=1), 2)
return pool_results, area_heights, area_widths
def embed(self, x):
"""Embedding function that takes discrete latent and returns embedding.
Args:
x: Input to the discretization bottleneck.
Returns:
Continuous embedding to be passed on to the decoder.
Raises:
ValueError: For unknown or missing arguments.
"""
shape_x = common_layers.shape_list(x)
x_flat = tf.reshape(x, [-1, 1])
c = self.int_to_bit(x_flat, num_bits=self.hparams.z_size, base=2)
shape = common_layers.shape_list(c)
new_shape = shape
new_shape.append(self.hparams.num_blocks)
new_shape.append(int(self.hparams.z_size / self.hparams.num_blocks))
c = tf.to_int32(tf.reshape(c, shape=new_shape))
h1_shape = shape_x
h1_shape.append(self.hparams.hidden_size)
h1 = tf.zeros(dtype=tf.float32, shape=h1_shape)
c_int = self.bit_to_int(c,
num_bits=int(self.hparams.z_size /
self.hparams.num_blocks),
base=2)
c_hot = tf.one_hot(c_int, depth=self.hparams.block_v_size, axis=-1)
c_hot_flat = tf.reshape(
c_hot,
shape=[-1, self.hparams.num_blocks, self.hparams.block_v_size])
h1 = tf.matmul(tf.transpose(c_hot_flat, perm=[1, 0, 2]), self.means)
h1 = tf.transpose(h1, perm=[1, 0, 2])
h1 = tf.reshape(h1, shape=h1_shape)
h1_shape[0] = self.hparams.batch_size
h2 = tf.layers.dense(tf.nn.relu(h1),
self.hparams.filter_size,
name="vch2")
res = tf.layers.dense(tf.nn.relu(h2),
self.hparams.hidden_size,
name="vcfin")
return res
def post_attention(self, token, x):
"""Called after self-attention. The memory can be updated here.
Args:
token: Data returned by pre_attention, which can be used to carry over
state related to the current memory operation.
x: a Tensor of data after self-attention and feed-forward
Returns:
a (possibly modified) version of the input x
"""
with tf.variable_scope(self.name + "/post_attention",
reuse=tf.AUTO_REUSE):
depth = common_layers.shape_list(x)[-1]
actual_batch_size = common_layers.shape_list(x)[0]
memory_output = tf.gather(token["retrieved_mem"],
tf.range(actual_batch_size))
output = tf.add(tf.layers.dense(x, depth, use_bias=False),
tf.layers.dense(memory_output, depth))
with tf.control_dependencies([output]):
with tf.control_dependencies(
[self.write(token["x"], token["access_logits"])]):
return tf.identity(output)
def _merge_beam_dim(tensor):
"""Reshapes first two dimensions in to single dimension.
Args:
tensor: Tensor to reshape of shape [A, B, ...]
Returns:
Reshaped tensor of shape [A*B, ...]
"""
shape = common_layers.shape_list(tensor)
shape[0] *= shape[1] # batch -> batch * beam_size
shape.pop(1) # Remove beam dim
return tf.reshape(tensor, shape)
def padded_accuracy(predictions,
labels,
weights_fn=common_layers.weights_nonzero):
"""Percentage of times that predictions matches labels on non-0s."""
# If the last dimension is 1 then we're using L1/L2 loss.
if common_layers.shape_list(predictions)[-1] == 1:
return rounding_accuracy(predictions, labels, weights_fn=weights_fn)
with tf.variable_scope("padded_accuracy", values=[predictions, labels]):
padded_predictions, padded_labels = common_layers.pad_with_zeros(
predictions, labels)
weights = weights_fn(padded_labels)
outputs = tf.to_int32(tf.argmax(padded_predictions, axis=-1))
padded_labels = tf.to_int32(padded_labels)
return to_float(tf.equal(outputs, padded_labels)), weights
def _gather(params, indices):
"""Fast gather using one_hot and batch matmul."""
if dtype != tf.float32:
params = to_float(params)
shape = common_layers.shape_list(params)
indices_shape = common_layers.shape_list(indices)
ndims = params.shape.ndims
# Adjust the shape of params to match one-hot indices, which is the
# requirement of Batch MatMul.
if ndims == 2:
params = tf.expand_dims(params, axis=-1)
if ndims > 3:
params = tf.reshape(params, [shape[0], shape[1], -1])
gather_result = tf.matmul(
tf.one_hot(indices, shape[1], dtype=params.dtype), params)
if ndims == 2:
gather_result = tf.squeeze(gather_result, axis=-1)
if ndims > 3:
shape[1] = indices_shape[1]
gather_result = tf.reshape(gather_result, shape)
if dtype != tf.float32:
gather_result = tf.cast(gather_result, dtype)
return gather_result
def _unmerge_beam_dim(tensor, batch_size, beam_size):
"""Reshapes first dimension back to [batch_size, beam_size].
Args:
tensor: Tensor to reshape of shape [batch_size*beam_size, ...]
batch_size: Tensor, original batch size.
beam_size: int, original beam size.
Returns:
Reshaped tensor of shape [batch_size, beam_size, ...]
"""
shape = common_layers.shape_list(tensor)
new_shape = [batch_size] + [beam_size] + shape[1:]
return tf.reshape(tensor, new_shape)
def _mix_tokens(p_sample, gold_targets, sampled_targets):
"""Interleave sampled and gold tokens randomly.
Args:
p_sample: float in [0, 1]. Probability a token will come from
'sampled_targets'. 0 means all-gold, 1 means all-sampled.
gold_targets: Tensor. Gold token IDs.
sampled_targets: Tensor. Sampled token IDs. Same shape as 'gold_targets'.
Returns:
Tensor of same shape as 'gold_targets' containing a mix of tokens from
'gold_targets' and 'sampled_targets'.
"""
targets_shape = common_layers.shape_list(sampled_targets)
return tf.where(tf.less(tf.random_uniform(targets_shape), p_sample),
sampled_targets, gold_targets)
def padded_accuracy_topk(predictions,
labels,
k,
weights_fn=common_layers.weights_nonzero):
"""Percentage of times that top-k predictions matches labels on non-0s."""
with tf.variable_scope("padded_accuracy_topk", values=[predictions, labels]):
padded_predictions, padded_labels = common_layers.pad_with_zeros(
predictions, labels)
weights = weights_fn(padded_labels)
effective_k = tf.minimum(k,
common_layers.shape_list(padded_predictions)[-1])
_, outputs = tf.nn.top_k(padded_predictions, k=effective_k)
outputs = tf.to_int32(outputs)
padded_labels = tf.to_int32(padded_labels)
padded_labels = tf.expand_dims(padded_labels, axis=-1)
padded_labels += tf.zeros_like(outputs) # Pad to same shape.
same = to_float(tf.equal(outputs, padded_labels))
same_topk = tf.reduce_sum(same, axis=-1)
return same_topk, weights
def sequence_edit_distance(predictions,
labels,
weights_fn=common_layers.weights_nonzero):
"""Average edit distance, ignoring padding 0s.
The score returned is the edit distance divided by the total length of
reference truth and the weight returned is the total length of the truth.
Args:
predictions: Tensor of shape [`batch_size`, `length`, 1, `num_classes`] and
type tf.float32 representing the logits, 0-padded.
labels: Tensor of shape [`batch_size`, `length`, 1, 1] and type tf.int32
representing the labels of same length as logits and 0-padded.
weights_fn: ignored. The weights returned are the total length of the ground
truth labels, excluding 0-paddings.
Returns:
(edit distance / reference length, reference length)
Raises:
ValueError: if weights_fn is not common_layers.weights_nonzero.
"""
if weights_fn is not common_layers.weights_nonzero:
raise ValueError("Only weights_nonzero can be used for this metric.")
with tf.variable_scope("edit_distance", values=[predictions, labels]):
# Transform logits into sequence classes by taking max at every step.
predictions = tf.to_int32(
tf.squeeze(tf.argmax(predictions, axis=-1), axis=(2, 3)))
nonzero_idx = tf.where(tf.not_equal(predictions, 0))
sparse_outputs = tf.SparseTensor(nonzero_idx,
tf.gather_nd(predictions, nonzero_idx),
tf.shape(predictions, out_type=tf.int64))
labels = tf.squeeze(labels, axis=(2, 3))
nonzero_idx = tf.where(tf.not_equal(labels, 0))
label_sparse_outputs = tf.SparseTensor(nonzero_idx,
tf.gather_nd(labels, nonzero_idx),
tf.shape(labels, out_type=tf.int64))
distance = tf.reduce_sum(
tf.edit_distance(sparse_outputs, label_sparse_outputs, normalize=False))
reference_length = to_float(common_layers.shape_list(nonzero_idx)[0])
return distance / reference_length, reference_length
def padded_neg_log_perplexity_with_masking(
predictions,
labels,
features,
weights_fn=None):
"""Average log-perplexity with custom targets_mask."""
del weights_fn
if "targets_mask" not in features:
raise ValueError("masked_neg_log_perplexity requires targets_mask feature")
# Features are 4 dimensional, so we need to reshape the targets_mask to match
# the shape of the labels. A lot of models rely on these features being 4D,
# so it's best to update the shape of the mask.
extended_targets_mask_shape = common_layers.shape_list(
features["targets_mask"])
extended_targets_mask_shape.extend([1, 1])
features["targets_mask"] = tf.reshape(features["targets_mask"],
shape=extended_targets_mask_shape)
mask_fn = lambda labels: features["targets_mask"]
return padded_neg_log_perplexity(predictions, labels, mask_fn)
def loss_fn(self, targets, logits):
"""Constructs loss dict.
Args:
targets: [batch_size, seq_len]
logits: [batch_size, seq_len, vocab_size]
Returns:
{str: Tensor of shape []}. Losses.
"""
batch_size, seq_len, vocab_size = common_layers.shape_list(logits)
targets = tf.reshape(targets, [batch_size, seq_len, 1, 1])
logits = tf.reshape(logits, [batch_size, seq_len, 1, 1, vocab_size])
features = copy.copy(self._features)
features["targets"] = targets
with tf.variable_scope(tf.get_variable_scope(), reuse=True):
losses = {
"training": self._t2tmodel.loss(logits, features),
}
return losses
def bit_to_int(self, x_bit, num_bits, base=2):
"""Turn x_bit representing numbers bitwise (lower-endian) to int tensor.
Args:
x_bit: Tensor containing numbers in a particular base to be
converted to
int.
num_bits: Number of bits in the representation.
base: Base of the representation.
Returns:
Integer representation of this number.
"""
x_l = tf.stop_gradient(tf.to_int32(tf.reshape(x_bit, [-1, num_bits])))
# pylint: disable=g-complex-comprehension
x_labels = [
x_l[:, i] * tf.to_int32(base)**tf.to_int32(i)
for i in range(num_bits)
]
res = sum(x_labels)
return tf.to_int32(
tf.reshape(res,
common_layers.shape_list(x_bit)[:-1]))
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,
layer_collection=None):
"""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
layer_collection: A tensorflow_kfac.LayerCollection. Only used by the
KFAC optimizer. Default is None.
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":
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,
layer_collection=layer_collection)
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_eval
if hparams.mode == tf.estimator.ModeKeys.TRAIN:
overhead = hparams.moe_overhead_train
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_eval
if hparams.mode == tf.estimator.ModeKeys.TRAIN:
overhead = hparams.moe_overhead_train
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 beam_search(symbols_to_logits_fn,
initial_ids,
beam_size,
decode_length,
vocab_size,
alpha,
states=None,
eos_id=EOS_ID,
stop_early=True,
use_tpu=False,
use_top_k_with_unique=True):
"""Beam search with length penalties.
Requires a function that can take the currently decoded symbols and return
the logits for the next symbol. The implementation is inspired by
https://arxiv.org/abs/1609.08144.
When running, the beam search steps can be visualized by using tfdbg to watch
the operations generating the output ids for each beam step. These operations
have the pattern:
(alive|finished)_topk_(seq,scores)
Operations marked `alive` represent the new beam sequences that will be
processed in the next step. Operations marked `finished` represent the
completed beam sequences, which may be padded with 0s if no beams finished.
Operations marked `seq` store the full beam sequence for the time step.
Operations marked `scores` store the sequence's final log scores.
The beam search steps will be processed sequentially in order, so when
capturing observed from these operations, tensors, clients can make
assumptions about which step is being recorded.
WARNING: Assumes 2nd dimension of tensors in `states` and not invariant, this
means that the shape of the 2nd dimension of these tensors will not be
available (i.e. set to None) inside symbols_to_logits_fn.
Args:
symbols_to_logits_fn: Interface to the model, to provide logits.
Shoud take [batch_size, decoded_ids] and return [batch_size, vocab_size]
initial_ids: Ids to start off the decoding, this will be the first thing
handed to symbols_to_logits_fn (after expanding to beam size)
[batch_size]
beam_size: Size of the beam.
decode_length: Number of steps to decode for.
vocab_size: Size of the vocab, must equal the size of the logits returned by
symbols_to_logits_fn
alpha: alpha for length penalty.
states: dict (possibly nested) of decoding states.
eos_id: ID for end of sentence.
stop_early: a boolean - stop once best sequence is provably determined.
use_tpu: A bool, whether to do beam search on TPU.
use_top_k_with_unique: bool, whether to use a fast (but decreased precision)
top_k during TPU beam search.
Returns:
Tuple of
(decoded beams [batch_size, beam_size, decode_length]
decoding probabilities [batch_size, beam_size])
"""
batch_size = common_layers.shape_list(initial_ids)[0]
# Assume initial_ids are prob 1.0
initial_log_probs = tf.constant([[0.] + [-INF] * (beam_size - 1)])
# Expand to beam_size (batch_size, beam_size)
alive_log_probs = tf.tile(initial_log_probs, [batch_size, 1])
# Expand each batch and state to beam_size
alive_seq = _expand_to_beam_size(initial_ids, beam_size)
alive_seq = tf.expand_dims(alive_seq, axis=2) # (batch_size, beam_size, 1)
if use_tpu:
alive_seq = tf.tile(alive_seq, [1, 1, decode_length + 1])
if states:
states = nest.map_structure(
lambda state: _expand_to_beam_size(state, beam_size), states)
else:
states = {}
# Finished will keep track of all the sequences that have finished so far
# Finished log probs will be negative infinity in the beginning
# finished_flags will keep track of booleans
finished_seq = tf.zeros(common_layers.shape_list(alive_seq), tf.int32)
# Setting the scores of the initial to negative infinity.
finished_scores = tf.ones([batch_size, beam_size]) * -INF
finished_flags = tf.zeros([batch_size, beam_size], tf.bool)
def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq,
curr_scores, curr_finished):
"""Given sequences and scores, will gather the top k=beam size sequences.
Args:
finished_seq: Current finished sequences.
[batch_size, beam_size, current_decoded_length]
finished_scores: scores for each of these sequences.
[batch_size, beam_size]
finished_flags: finished bools for each of these sequences.
[batch_size, beam_size]
curr_seq: current topk sequence that has been grown by one position.
[batch_size, beam_size, current_decoded_length]
curr_scores: scores for each of these sequences. [batch_size, beam_size]
curr_finished: Finished flags for each of these sequences.
[batch_size, beam_size]
Returns:
Tuple of
(Topk sequences based on scores,
log probs of these sequences,
Finished flags of these sequences)
"""
if not use_tpu:
# First append a column of 0'ids to finished to make the same length with
# finished scores
finished_seq = tf.concat(
[finished_seq,
tf.zeros([batch_size, beam_size, 1], tf.int32)],
axis=2)
# Set the scores of the unfinished seq in curr_seq to large negative
# values
curr_scores += (1. - to_float(curr_finished)) * -INF
# concatenating the sequences and scores along beam axis
curr_finished_seq = tf.concat([finished_seq, curr_seq], axis=1)
curr_finished_scores = tf.concat([finished_scores, curr_scores],
axis=1)
curr_finished_flags = tf.concat([finished_flags, curr_finished],
axis=1)
return compute_topk_scores_and_seq(
curr_finished_seq,
curr_finished_scores,
curr_finished_scores,
curr_finished_flags,
beam_size,
batch_size,
"grow_finished",
use_tpu=use_tpu,
use_top_k_with_unique=use_top_k_with_unique)
def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished,
states):
"""Given sequences and scores, will gather the top k=beam size sequences.
Args:
curr_seq: current topk sequence that has been grown by one position.
[batch_size, beam_size, i+1]
curr_scores: scores for each of these sequences. [batch_size, beam_size]
curr_log_probs: log probs for each of these sequences.
[batch_size, beam_size]
curr_finished: Finished flags for each of these sequences.
[batch_size, beam_size]
states: dict (possibly nested) of decoding states.
Returns:
Tuple of
(Topk sequences based on scores,
log probs of these sequences,
Finished flags of these sequences)
"""
# Set the scores of the finished seq in curr_seq to large negative
# values
curr_scores += to_float(curr_finished) * -INF
return compute_topk_scores_and_seq(curr_seq,
curr_scores,
curr_log_probs,
curr_finished,
beam_size,
batch_size,
"grow_alive",
states,
use_tpu=use_tpu)
def grow_topk(i, alive_seq, alive_log_probs, states):
r"""Inner beam search loop.
This function takes the current alive sequences, and grows them to topk
sequences where k = 2*beam. We use 2*beam because, we could have beam_size
number of sequences that might hit <EOS> and there will be no alive
sequences to continue. With 2*beam_size, this will not happen. This relies
on the assumption the vocab size is > beam size. If this is true, we'll
have at least beam_size non <EOS> extensions if we extract the next top
2*beam words.
Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to
https://arxiv.org/abs/1609.08144.
Args:
i: loop index
alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1]
alive_log_probs: probabilities of these sequences. [batch_size, beam_size]
states: dict (possibly nested) of decoding states.
Returns:
Tuple of
(Topk sequences extended by the next word,
The log probs of these sequences,
The scores with length penalty of these sequences,
Flags indicating which of these sequences have finished decoding,
dict of transformed decoding states)
"""
# Get the logits for all the possible next symbols
if use_tpu and states:
flat_ids = tf.reshape(
tf.slice(alive_seq, [0, 0, i], [batch_size, beam_size, 1]),
[batch_size * beam_size, -1])
else:
flat_ids = tf.reshape(alive_seq, [batch_size * beam_size, -1])
# (batch_size * beam_size, decoded_length)
if states:
flat_states = nest.map_structure(_merge_beam_dim, states)
flat_logits, flat_states = symbols_to_logits_fn(
flat_ids, i, flat_states)
states = nest.map_structure(
lambda t: _unmerge_beam_dim(t, batch_size, beam_size),
flat_states)
elif use_tpu:
flat_logits = symbols_to_logits_fn(flat_ids, i)
else:
flat_logits = symbols_to_logits_fn(flat_ids)
logits = tf.reshape(flat_logits, [batch_size, beam_size, -1])
# Convert logits to normalized log probs
candidate_log_probs = common_layers.log_prob_from_logits(logits)
# Multiply the probabilities by the current probabilities of the beam.
# (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1)
log_probs = candidate_log_probs + tf.expand_dims(alive_log_probs,
axis=2)
length_penalty = tf.pow(((5. + to_float(i + 1)) / 6.), alpha)
curr_scores = log_probs / length_penalty
# Flatten out (beam_size, vocab_size) probs in to a list of possibilities
flat_curr_scores = tf.reshape(curr_scores,
[-1, beam_size * vocab_size])
if use_tpu and use_top_k_with_unique:
topk_scores, topk_ids = top_k_with_unique(flat_curr_scores,
k=beam_size * 2)
else:
topk_scores, topk_ids = tf.nn.top_k(flat_curr_scores,
k=beam_size * 2)
# Recovering the log probs because we will need to send them back
topk_log_probs = topk_scores * length_penalty
# Work out what beam the top probs are in.
topk_beam_index = topk_ids // vocab_size
topk_ids %= vocab_size # Unflatten the ids
if not use_tpu:
# The next three steps are to create coordinates for tf.gather_nd to pull
# out the correct sequences from id's that we need to grow.
# We will also use the coordinates to gather the booleans of the beam
# items that survived.
batch_pos = compute_batch_indices(batch_size, beam_size * 2)
# top beams will give us the actual coordinates to do the gather.
# stacking will create a tensor of dimension batch * beam * 2, where the
# last dimension contains the i,j gathering coordinates.
topk_coordinates = tf.stack([batch_pos, topk_beam_index], axis=2)
# Gather up the most probable 2*beams both for the ids and
# finished_in_alive bools
topk_seq = tf.gather_nd(alive_seq, topk_coordinates)
if states:
states = nest.map_structure(
lambda state: tf.gather_nd(state, topk_coordinates),
states)
# Append the most probable alive
topk_seq = tf.concat(
[topk_seq, tf.expand_dims(topk_ids, axis=2)], axis=2)
else:
# Gather up the most probable 2*beams both for the ids and
# finished_in_alive bools
topk_seq = fast_tpu_gather(alive_seq, topk_beam_index)
if states:
states = nest.map_structure(
lambda state: fast_tpu_gather(state, topk_beam_index),
states)
# Update the most probable alive
topk_seq = tf.transpose(topk_seq, perm=[2, 0, 1])
topk_seq = inplace_ops.alias_inplace_update(
topk_seq, i + 1, topk_ids)
topk_seq = tf.transpose(topk_seq, perm=[1, 2, 0])
topk_finished = tf.equal(topk_ids, eos_id)
return topk_seq, topk_log_probs, topk_scores, topk_finished, states
def inner_loop(i, alive_seq, alive_log_probs, finished_seq,
finished_scores, finished_flags, states):
"""Inner beam search loop.
There are three groups of tensors, alive, finished, and topk.
The alive group contains information about the current alive sequences
The topk group contains information about alive + topk current decoded words
the finished group contains information about finished sentences, that is,
the ones that have decoded to <EOS>. These are what we return.
The general beam search algorithm is as follows:
While we haven't terminated (pls look at termination condition)
1. Grow the current alive to get beam*2 topk sequences
2. Among the topk, keep the top beam_size ones that haven't reached EOS
into alive
3. Among the topk, keep the top beam_size ones have reached EOS into
finished
Repeat
To make things simple with using fixed size tensors, we will end
up inserting unfinished sequences into finished in the beginning. To stop
that we add -ve INF to the score of the unfinished sequence so that when a
true finished sequence does appear, it will have a higher score than all the
unfinished ones.
Args:
i: loop index
alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1]
alive_log_probs: probabilities of the beams. [batch_size, beam_size]
finished_seq: Current finished sequences.
[batch_size, beam_size, i+1]
finished_scores: scores for each of these sequences.
[batch_size, beam_size]
finished_flags: finished bools for each of these sequences.
[batch_size, beam_size]
states: dict (possibly nested) of decoding states.
Returns:
Tuple of
(Incremented loop index
New alive sequences,
Log probs of the alive sequences,
New finished sequences,
Scores of the new finished sequences,
Flags indicating which sequence in finished as reached EOS,
dict of final decoding states)
"""
# Each inner loop, we carry out three steps:
# 1. Get the current topk items.
# 2. Extract the ones that have finished and haven't finished
# 3. Recompute the contents of finished based on scores.
topk_seq, topk_log_probs, topk_scores, topk_finished, states = grow_topk(
i, alive_seq, alive_log_probs, states)
alive_seq, alive_log_probs, _, states = grow_alive(
topk_seq, topk_scores, topk_log_probs, topk_finished, states)
finished_seq, finished_scores, finished_flags, _ = grow_finished(
finished_seq, finished_scores, finished_flags, topk_seq,
topk_scores, topk_finished)
return (i + 1, alive_seq, alive_log_probs, finished_seq,
finished_scores, finished_flags, states)
def _is_not_finished(i, unused_alive_seq, alive_log_probs,
unused_finished_seq, finished_scores,
unused_finished_in_finished, unused_states):
"""Checking termination condition.
We terminate when we decoded up to decode_length or the lowest scoring item
in finished has a greater score that the highest prob item in alive divided
by the max length penalty
Args:
i: loop index
alive_log_probs: probabilities of the beams. [batch_size, beam_size]
finished_scores: scores for each of these sequences.
[batch_size, beam_size]
Returns:
Bool.
"""
max_length_penalty = tf.pow(((5. + to_float(decode_length)) / 6.),
alpha)
# The best possible score of the most likely alive sequence.
lower_bound_alive_scores = alive_log_probs[:, 0] / max_length_penalty
if not stop_early:
# by considering the min score (in the top N beams) we ensure that
# the decoder will keep decoding until there is at least one beam
# (in the top N) that can be improved (w.r.t. the alive beams).
# any unfinished beam will have score -INF - thus the min
# will always be -INF if there is at least one unfinished beam -
# which means the bound_is_met condition cannot be true in this case.
lowest_score_of_finished_in_finished = tf.reduce_min(
finished_scores)
else:
# by taking the max score we only care about the first beam;
# as soon as this first beam cannot be beaten from the alive beams
# the beam decoder can stop.
# similarly to the above, if the top beam is not completed, its
# finished_score is -INF, thus it will not activate the
# bound_is_met condition. (i.e., decoder will keep going on).
# note we need to find the max for every sequence eparately - so, we need
# to keep the batch dimension (see axis=1)
lowest_score_of_finished_in_finished = tf.reduce_max(
finished_scores, axis=1)
bound_is_met = tf.reduce_all(
tf.greater(lowest_score_of_finished_in_finished,
lower_bound_alive_scores))
return tf.logical_and(tf.less(i, decode_length),
tf.logical_not(bound_is_met))
inner_shape = tf.TensorShape([None, None, None])
if use_tpu:
inner_shape = tf.TensorShape(
[batch_size, beam_size, decode_length + 1])
if use_tpu:
state_struc = nest.map_structure(lambda state: state.get_shape(),
states)
else:
state_struc = nest.map_structure(get_state_shape_invariants, states)
(_, alive_seq, alive_log_probs, finished_seq, finished_scores,
finished_flags, states) = tf.while_loop(
_is_not_finished,
inner_loop, [
tf.constant(0), alive_seq, alive_log_probs, finished_seq,
finished_scores, finished_flags, states
],
shape_invariants=[
tf.TensorShape([]), inner_shape,
alive_log_probs.get_shape(), inner_shape,
finished_scores.get_shape(),
finished_flags.get_shape(), state_struc
],
parallel_iterations=1,
back_prop=False)
alive_seq.set_shape((None, beam_size, None))
finished_seq.set_shape((None, beam_size, None))
# Accounting for corner case: It's possible that no sequence in alive for a
# particular batch item ever reached EOS. In that case, we should just copy
# the contents of alive for that batch item. tf.reduce_any(finished_flags, 1)
# if 0, means that no sequence for that batch index had reached EOS. We need
# to do the same for the scores as well.
finished_seq = tf.where(tf.reduce_any(finished_flags, 1), finished_seq,
alive_seq)
finished_scores = tf.where(tf.reduce_any(finished_flags, 1),
finished_scores, alive_log_probs)
return finished_seq, finished_scores, states
def sequential_scheduled_sampling(infer_fn, mix_fn, loss_fn, targets):
"""Scheduled Sampling.
Args:
infer_fn: Function. Computes logits for all timesteps.
mix_fn: Function. Mixes gold and sample tokens.
loss_fn: Function. Computes loss between gold tokens and logits.
targets: Tensor of shape [batch_size, seq_len]. Gold tokens.
Returns:
ss_tokens: Tensor of shape [batch_size, seq_len]. Scheduled sampling tokens.
ss_logits: Tensor of shape [batch_size, seq_len, vocab_size]. Logits for
next token when conditioning on ss_tokens.
losses_dict: {str: scalar Tensor}. Losses to optimize.
"""
targets_shape = common_layers.shape_list(targets)
batch_size = targets_shape[0]
seq_len = targets_shape[1]
if not targets.shape.is_fully_defined():
# TODO(duckworthd): When running on GPU, I get the following error. Solve
# it to enable use on other devices.
#
# Cannot use 'Identity_186' as input to
# 'transformer/parallel_0_7/transformer/transformer/symbol_modality_16282_512/shared/convert_gradient_to_tensor_HBc3xYw22Mw'
# because 'Identity_186' is in a while loop.
raise ValueError(
"The following code only works on TPU. As targets.shape isn't fully "
"defined, I am assuming you are using a different device.")
def cond_fn(i, ss_tokens):
"""True if i < seq_len."""
_ = ss_tokens
return i < seq_len
def body_fn(i, ss_tokens):
"""Constructs conditioning tokens for scheduled sampling."""
# next_token_logits depends on timesteps 0...i-1.
#
# [batch_size, seq_len] -> [batch_size, seq_len, vocab_size]
ss_tokens_logits = infer_fn(ss_tokens)
# Same as 'next_token_logits = ss_tokens_logits[:, i, :]'.
vocab_size = common_layers.shape_list(ss_tokens_logits)[2]
next_token_logits = tf.slice(ss_tokens_logits,
begin=[0, i, 0],
size=[batch_size, 1, vocab_size])
next_token_logits = tf.squeeze(next_token_logits, axis=[1])
# [batch_size, vocab_size] -> [batch_size]
sampled_next_tokens = _sample_next_tokens(next_token_logits)
# Same as 'gold_next_tokens = targets[:, i]'.
gold_next_tokens = tf.slice(targets,
begin=[0, i],
size=[batch_size, 1])
gold_next_tokens = tf.squeeze(gold_next_tokens, axis=[1])
next_tokens = mix_fn(gold_next_tokens, sampled_next_tokens)
ss_tokens = _update_timestep(ss_tokens, timestep=i, values=next_tokens)
return i + 1, tf.stop_gradient(ss_tokens)
# tf.while_loop() over all timesteps. Generate scheduled sampling tokens.
i = 0
ss_tokens = tf.zeros([batch_size, seq_len], dtype=tf.int32)
i, ss_tokens = tf.while_loop(cond_fn, body_fn, [i, ss_tokens])
ss_logits = infer_fn(ss_tokens)
return ss_tokens, ss_logits, loss_fn(targets, ss_logits)
def _compute_sum_image(features,
max_area_width,
max_area_height=1,
height=1,
name=None):
"""Computes area sums for features.
Args:
features: a Tensor in a shape of [batch_size, height * width, depth].
max_area_width: the max width allowed for an area.
max_area_height: the max height allowed for an area.
height: the height of the image.
name: the namescope.
Returns:
sum_image: A Tensor of shape [batch_size, num_areas, depth]
area_heights: A Tensor of shape [batch_size, num_areas, 1]
area_widths: A Tensor of shape [batch_size, num_areas, 1]
"""
with tf.name_scope(name, default_name="compute_sum_image"):
feature_shape = common_layers.shape_list(features)
batch_size = feature_shape[0]
length = feature_shape[-2]
depth = feature_shape[-1]
width = length // height
features_2d = tf.reshape(features, [batch_size, height, width, depth])
width_cum = tf.cumsum(features_2d, axis=-2, name="compute_integral_h")
integral_image = tf.cumsum(width_cum,
axis=-3,
name="compute_integral_v")
padded_image = tf.pad(integral_image, [[0, 0], [1, 0], [1, 0], [0, 0]],
constant_values=0)
height_list = []
width_list = []
dst_images = []
src_images_diag = []
src_images_h = []
src_images_v = []
size_tensor = tf.ones_like(padded_image[:, :, :, 0], dtype=tf.int32)
for area_height in range(max_area_height):
for area_width in range(max_area_width):
dst_images.append(
tf.reshape(
padded_image[:, area_height + 1:, area_width + 1:, :],
[batch_size, -1, depth]))
src_images_diag.append(
tf.reshape(
padded_image[:, :-area_height - 1, :-area_width -
1, :], [batch_size, -1, depth]))
src_images_h.append(
tf.reshape(
padded_image[:, area_height + 1:, :-area_width - 1, :],
[batch_size, -1, depth]))
src_images_v.append(
tf.reshape(
padded_image[:, :-area_height - 1, area_width + 1:, :],
[batch_size, -1, depth]))
height_list.append(
tf.reshape(
size_tensor[:, area_height + 1:, area_width + 1:] *\
(area_height + 1), [batch_size, -1]))
width_list.append(
tf.reshape(
size_tensor[:, area_height + 1:, area_width + 1:] *\
(area_width + 1), [batch_size, -1]))
sum_image = tf.subtract(
tf.concat(dst_images, axis=1) + tf.concat(src_images_diag, axis=1),
tf.concat(src_images_v, axis=1) + tf.concat(src_images_h, axis=1))
area_heights = tf.expand_dims(tf.concat(height_list, axis=1), 2)
area_widths = tf.expand_dims(tf.concat(width_list, axis=1), 2)
return sum_image, area_heights, area_widths
def compute_area_key(features,
max_area_width,
max_area_height=1,
height=1,
mode="mean",
training=True,
name=None):
"""Computes the key for each area.
Args:
features: a Tensor in a shape of [batch_size, height * width, depth].
max_area_width: the max width allowed for an area.
max_area_height: the max height allowed for an area.
height: the height of the image.
mode: whether to combine different area features or only use
the vector mean of each area, which can be "mean", "concat", "sum",
"sample_concat", and "sample_sum".
training: indicating if it is in the training mode.
name: the name for setting the variable scope.
Returns:
area_key: a Tensor in the shape of [batch_size, num_areas, depth]
"""
tf.logging.info("area_attention mode=%s", mode)
area_mean, area_std, _, area_heights, area_widths =\
compute_area_features(features, max_area_width=max_area_width,
max_area_height=max_area_height, height=height)
if mode == "mean":
return area_mean
elif mode == "max":
area_max, _, _ = basic_pool(features,
max_area_width=max_area_width,
max_area_height=max_area_height,
height=height)
return area_max
elif mode == "sample":
if training:
area_mean += (area_std * tf.random_normal(tf.shape(area_std)))
return area_mean
with tf.variable_scope(
name,
default_name="combine_area_features",
values=[area_mean, area_std, area_heights, area_widths]):
depth = common_layers.shape_list(area_mean)[-1]
height_embed = tf.nn.embedding_lookup(params=tf.get_variable(
"area_height_emb", [max_area_height, depth // 2]),
ids=area_heights[:, :, 0] - 1)
width_embed = tf.nn.embedding_lookup(params=tf.get_variable(
"area_width_emb", [max_area_width, depth // 2]),
ids=area_widths[:, :, 0] - 1)
size_embed = tf.concat([height_embed, width_embed], -1)
if mode == "concat":
feature_concat = tf.concat([area_mean, area_std, size_embed], -1)
elif mode == "max_concat":
area_max, _, _ = basic_pool(features,
max_area_width=max_area_width,
max_area_height=max_area_height,
height=height)
feature_concat = tf.concat([area_max, size_embed], -1)
elif mode == "sum":
feature_concat = size_embed + area_mean + area_std
elif mode == "sample_concat":
if training:
area_mean += (area_std * tf.random_normal(tf.shape(area_std)))
feature_concat = tf.concat([area_mean, size_embed], -1)
elif mode == "sample_sum":
if training:
area_mean += (area_std * tf.random_normal(tf.shape(area_std)))
feature_concat = area_mean + size_embed
else:
raise ValueError("Unsupported area key mode=%s" % mode)
feature_hidden = tf.layers.dense(inputs=feature_concat,
units=depth,
activation=tf.nn.relu)
area_key = tf.layers.dense(feature_hidden, units=depth)
return area_key
def dot_product_area_attention(q,
k,
v,
bias,
dropout_rate=0.0,
image_shapes=None,
name=None,
attention_image_summary=None,
save_weights_to=None,
dropout_broadcast_dims=None,
max_area_width=1,
max_area_height=1,
memory_height=1,
area_key_mode="mean",
area_value_mode="sum",
top_k_areas=0,
area_temperature=1.0,
training=True):
"""Dot-product area attention.
Args:
q: Tensor with shape [..., length_q, depth_k].
k: Tensor with shape [..., length_kv, depth_k]. Leading dimensions must
match with q.
v: Tensor with shape [..., length_kv, depth_v] Leading dimensions must
match with q.
bias: bias Tensor (see attention_bias())
dropout_rate: a float.
image_shapes: optional tuple of integer scalars.
see comments for attention_image_summary()
name: an optional string
attention_image_summary: the callback for making image summary of attention.
save_weights_to: an optional dictionary to capture attention weights
for visualization; the weights tensor will be appended there under
a string key created from the variable scope (including name).
dropout_broadcast_dims: an optional list of integers less than rank of q.
Specifies in which dimensions to broadcast the dropout decisions.
max_area_width: the max width allowed for an area.
max_area_height: the max height allowed for an area.
memory_height: the height of the memory.
area_key_mode: the mode for computing area keys, which can be "mean",
"concat", "sum", "sample_concat", and "sample_sum".
area_value_mode: the mode for computing area values, which can be either
"mean", or "sum".
top_k_areas: Use the top key areas for attention.
area_temperature: the temperature for attention softmax.
training: indicating if it is in the training mode.
Returns:
Tensor with shape [..., length_q, depth_v].
"""
tf.logging.info(
"dot_product_area_attention: "
"area_h=%d, area_w=%d, mem_h=%d, "
"area_key_mode=%s, area_value_mode=%s, "
"area_temperature=%f", max_area_height, max_area_width, memory_height,
area_key_mode, area_value_mode, area_temperature)
with tf.variable_scope(name,
default_name="dot_product_area_attention",
values=[q, k, v]) as scope:
mem_shape = common_layers.shape_list(k)
batch_size = mem_shape[0]
head_size = mem_shape[1]
length = mem_shape[2]
depth = mem_shape[3]
k_area = compute_area_key(tf.reshape(k, [-1, length, depth]),
max_area_width=max_area_width,
max_area_height=max_area_height,
height=memory_height,
mode=area_key_mode,
training=training)
if area_value_mode == "mean":
v_area, _, _, _, _ = compute_area_features(
tf.reshape(v, [-1, length, depth]),
max_area_width=max_area_width,
max_area_height=max_area_height,
height=memory_height)
elif area_value_mode == "max":
v_area, _, _ = basic_pool(tf.reshape(v, [-1, length, depth]),
max_area_width=max_area_width,
max_area_height=max_area_height,
height=memory_height,
fn=tf.reduce_max)
elif area_value_mode == "sum":
_, _, v_area, _, _ = compute_area_features(
tf.reshape(v, [-1, length, depth]),
max_area_width=max_area_width,
max_area_height=max_area_height,
height=memory_height)
else:
raise ValueError("Unsupported area value mode=%s" %
area_value_mode)
k = tf.reshape(k_area, [batch_size, head_size, -1, depth])
v = tf.reshape(v_area, [batch_size, head_size, -1, depth])
logits = tf.matmul(q, k,
transpose_b=True) # [..., length_q, length_kv]
if bias is not None:
bias = common_layers.cast_like(bias, logits)
with tf.name_scope("compute_area_att_bias", values=[bias]):
bias_shape = common_layers.shape_list(bias)
mem_length = bias_shape[-1]
bias_values = tf.reshape(
tf.cast(tf.less(bias, -1), tf.float32),
[-1, mem_length, 1])
_, _, padding_sum, _, _ = compute_area_features(
bias_values,
max_area_width=max_area_width,
max_area_height=max_area_height,
height=memory_height)
bias = tf.where(tf.cast(tf.to_int32(padding_sum), tf.bool),
tf.fill(tf.shape(padding_sum), -np.inf),
tf.zeros_like(padding_sum, dtype=tf.float32))
bias = tf.reshape(
bias, [bias_shape[0], bias_shape[1], bias_shape[2], -1])
logits += bias
logits = logits / area_temperature
weights = tf.nn.softmax(logits, name="attention_weights")
if top_k_areas > 0:
tf.logging.info("area_attention top_k_areas=%d", top_k_areas)
top_k = tf.minimum(
common_layers.shape_list(weights)[-1], top_k_areas)
top_weights, _ = tf.nn.top_k(weights, k=top_k)
min_values = tf.reduce_min(top_weights, -1, keepdims=True)
weights = tf.where(tf.greater_equal(weights, min_values), weights,
tf.zeros_like(weights))
weights = tf.div(weights, tf.reduce_sum(weights, -1,
keepdims=True))
if save_weights_to is not None:
save_weights_to[scope.name] = weights
save_weights_to[scope.name + "/logits"] = logits
# Drop out attention links for each head.
weights = common_layers.dropout_with_broadcast_dims(
weights, 1.0 - dropout_rate, broadcast_dims=dropout_broadcast_dims)
if common_layers.should_generate_summaries(
) and attention_image_summary:
attention_image_summary(weights, image_shapes)
return tf.matmul(weights, v)
def discrete_bottleneck(self, x):
"""Discretization bottleneck for latent variables.
Args:
x: Input to the discretization bottleneck.
Returns:
Embedding to pass to the decoder, discrete latent, loss, and the
embedding
function.
Raises:
ValueError: If projection_tensors is None for reshape_method
project, or
ema_count or ema_means is None if we are using ema, or unknown
args.
"""
x_reshaped = self.slice_hidden(x)
x_means_hot = []
x_means = 0
loss = 0
x_means_hot, x_means, q_loss, e_loss = self.embedding_lookup(
x_reshaped, self.means)
if self.hparams.ema:
tf.logging.info("Using EMA with beta = {}".format(
self.hparams.beta))
updated_ema_count = \
moving_averages.assign_moving_average(
self.ema_count,
tf.reduce_sum(
tf.reshape(
x_means_hot,
shape=[-1, self.hparams.num_blocks,
self.hparams.block_v_size]),
axis=0),
self.hparams.decay,
zero_debias=False)
dw = tf.matmul(tf.transpose(x_means_hot, perm=[1, 2, 0]),
tf.transpose(x_reshaped, perm=[1, 0, 2]))
updated_ema_means = \
moving_averages.assign_moving_average(
self.ema_means, dw, self.hparams.decay,
zero_debias=False)
n = tf.reduce_sum(updated_ema_count, axis=-1, keep_dims=True)
updated_ema_count = (
(updated_ema_count + self.hparams.epsilon) /
(n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
updated_ema_means = updated_ema_means / tf.expand_dims(
updated_ema_count, axis=-1)
with tf.control_dependencies([e_loss]):
update_means = tf.assign(self.means, updated_ema_means)
with tf.control_dependencies([update_means]):
loss += self.hparams.beta * e_loss
else:
# Use a gradient based loss for learning the cluster centers
loss += q_loss + self.hparams.beta * e_loss
# Get the discrete latent representation
x_means_idx = tf.argmax(x_means_hot, axis=-1)
# Get the binary representation
num_bits = int(self.hparams.z_size // self.hparams.num_blocks)
x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
x_discrete = self.bit_to_int(tf.to_int32(x_means_bits),
num_bits=self.hparams.z_size,
base=2)
# Reshape x_discrete
shape_x = common_layers.shape_list(x)
shape_discrete = shape_x[:-1]
x_discrete = tf.reshape(x_discrete, shape_discrete)
x_means = tf.reshape(x_means, shape=shape_x)
h1 = x + tf.stop_gradient(x_means - x)
h2 = tf.layers.dense(tf.nn.relu(h1),
self.hparams.filter_size,
name="vch2")
res = tf.layers.dense(tf.nn.relu(h2),
self.hparams.hidden_size,
name="vcfin")
embed_fn = partial(self.embed)
return {
"dense": res,
"discrete": x_discrete,
"loss": loss,
"embed": embed_fn
}
def class_label_targets_bottom(x, model_hparams, vocab_size):
with tf.variable_scope("class_label_modality_%d_%d" %
(vocab_size, model_hparams.hidden_size)):
return tf.zeros(
[common_layers.shape_list(x)[0], 1, 1, model_hparams.hidden_size])
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)
