def _Moments(inputs, mask, enable_cross_replica_sum_on_tpu=False):
"""Computes mean and variance over the valid data points in inputs."""
inputs = py_utils.with_dependencies([
py_utils.assert_equal(tf.rank(inputs), tf.rank(mask)),
py_utils.assert_greater_equal(mask, tf.zeros_like(mask)),
], inputs)
rank = tf.rank(mask)
reduce_over_dims = tf.range(0, rank - 1)
sum_v = tf.reduce_sum(inputs * tf.cast(mask, inputs.dtype),
reduce_over_dims)
count_v = tf.reduce_sum(mask, reduce_over_dims)
# Input shape is guaranteed to be a multiple of mask shape because the
# inputs * mask op above was successfully broadcasted.
mask_multiplier = tf.shape(inputs)[:-1] // tf.shape(mask)[:-1]
count_v *= tf.cast(tf.reduce_prod(mask_multiplier), count_v.dtype)
if py_utils.use_tpu() and enable_cross_replica_sum_on_tpu:
sum_v = tf.tpu.cross_replica_sum(sum_v)
count_v = tf.tpu.cross_replica_sum(count_v)
count_v = tf.maximum(count_v, 1.0)
mean = sum_v / count_v
sum_vv = tf.reduce_sum((inputs - mean) * (inputs - mean) * mask,
reduce_over_dims)
if py_utils.use_tpu() and enable_cross_replica_sum_on_tpu:
sum_vv = tf.tpu.cross_replica_sum(sum_vv)
variance = py_utils.with_dependencies([
py_utils.assert_greater_equal(sum_vv, tf.zeros_like(sum_vv)),
], sum_vv / count_v)
return mean, variance
def _DerivePaddingsAndIds(src_ids, tgt_labels):
"""tgt_ids is tgt_labels shifted right by one, with a SOS ID prepended."""
tgt_ids = tf.concat([[p.sos_id], tgt_labels[:-1]], axis=0)
src_paddings = tf.zeros(tf.shape(src_ids), dtype=tf.float32)
tgt_paddings = tf.zeros(tf.shape(tgt_ids), dtype=tf.float32)
tgt_weights = tf.ones(tf.shape(tgt_ids), dtype=tf.float32)
bucket_key = tf.cast(
tf.maximum(tf.reduce_sum(1.0 - src_paddings),
tf.reduce_sum(1.0 - tgt_paddings)), tf.int32)
return src_paddings, tgt_ids, tgt_paddings, tgt_weights, bucket_key
def SequenceConcat(x, x_paddings, y, y_paddings, pad=0):
"""Concats sequence `x` with sequence `y`.
This function is length aware (based off the paddings).
Args:
x: A sequence of tokens of shape [batch_size, x_len_max].
x_paddings: The paddings of `x`.
y: A sequence of tokens of shape [batch_size, y_len_max].
y_paddings: The paddings of `y`.
pad: The <pad> token to fill the concatenated sequence (of type integer).
Returns:
A tuple.
- Concatenation of `x` and `y` of shape
[batch_size, x_len_max + y_len_max].
- Paddings of the concatenation of shape
[batch_size, x_len_max + y_len_max].
"""
# Get the length (w/ eos).
x_len = tf.cast(tf.round(tf.reduce_sum(1 - x_paddings, 1)), tf.int32)
y_len = tf.cast(tf.round(tf.reduce_sum(1 - y_paddings, 1)), tf.int32)
batch_size = py_utils.GetShape(x)[0]
y_len_max = py_utils.GetShape(y)[1]
# Pad `x` with necessary <pad>.
x = tf.concat([x, tf.fill(py_utils.GetShape(y), pad)], 1)
# Replace all <pad> with 0.
x = tf.where(tf.not_equal(x, pad), x, tf.fill(py_utils.GetShape(x), 0))
# Compute the write indices of `y` in `xy`.
indices = tf.stack([
tf.tile(tf.expand_dims(tf.range(batch_size), 1), [1, y_len_max]),
(tf.tile(tf.expand_dims(tf.range(y_len_max), 0), [batch_size, 1]) +
tf.expand_dims(x_len, 1)),
], 2)
xy = x + tf.scatter_nd(indices, y, py_utils.GetShape(x))
# We need to remap all <pad> to `pad`.
xy = tf.where(
tf.less(tf.expand_dims(tf.range(py_utils.GetShape(xy)[1]), 0),
tf.expand_dims(x_len + y_len, 1)), xy,
tf.fill(py_utils.GetShape(xy), pad))
xy_paddings = 1 - tf.sequence_mask(x_len + y_len,
py_utils.GetShape(xy)[1],
x_paddings.dtype)
return xy, xy_paddings
def SequenceAppendToken(x, x_paddings, token, extend=False):
"""Appends <token> to sequence `x`.
Args:
x: A sequence of tokens of shape [batch_size, x_len_max].
x_paddings: The paddings of `x`.
token: The token to append (of type integer).
extend: Whether to extend `x` along the length dimension, this must be true
for any sequence length in `x` that is `x_len_max` or else an invalid
sequence will be emitted.
Returns:
A tuple.
- The new sequence, Tensor of shape [batch_size, x_len_max].
- The new paddings, Tensor of shape [batch_size, x_len_max].
"""
batch_size = py_utils.GetShape(x)[0]
x_len = tf.cast(tf.round(tf.reduce_sum(1 - x_paddings, 1)), tf.int32)
if extend:
x = tf.pad(x, [[0, 0], [0, 1]])
# Mask all invalid entries of `x` to 0.
x *= tf.sequence_mask(x_len, py_utils.GetShape(x)[1], x.dtype)
# Append the <token> based on `x_len`.
x += tf.scatter_nd(tf.stack([tf.range(batch_size), x_len], axis=1),
tf.cast(tf.fill([batch_size], token), x.dtype),
py_utils.GetShape(x))
x_paddings = 1 - tf.sequence_mask(x_len + 1,
py_utils.GetShape(x)[1],
x_paddings.dtype)
return x, x_paddings
def _ProcessBeamSearchDecodeOut(self, input_batch, encoder_outputs,
decoder_outs):
self.r1_shape = decoder_outs[0]
self.r2_shape = decoder_outs[1]
self.r3_shape = decoder_outs[2]
tf.logging.info('r1_shape: %s', self.r1_shape)
tf.logging.info('r2_shape: %s', self.r2_shape)
tf.logging.info('r3_shape: %s', self.r3_shape)
hyps = decoder_outs[3]
prev_hyps = decoder_outs[4]
done_hyps = decoder_outs[5]
scores = decoder_outs[6]
atten_probs = decoder_outs[7]
eos_scores = decoder_outs[8]
eos_atten_probs = decoder_outs[9]
source_seq_lengths = decoder_outs[10]
tlen = tf.cast(
tf.round(tf.reduce_sum(1.0 - input_batch.tgt.paddings, 1) - 1.0),
tf.int32)
ret_dict = {
'target_ids': input_batch.tgt.ids[:, 1:],
'eval_weight': input_batch.eval_weight,
'tlen': tlen,
'hyps': hyps,
'prev_hyps': prev_hyps,
'done_hyps': done_hyps,
'scores': scores,
'atten_probs': atten_probs,
'eos_scores': eos_scores,
'eos_atten_probs': eos_atten_probs,
'source_seq_lengths': source_seq_lengths,
}
return ret_dict
def EinsumBxycBzxBzyc(self, a, b, name=None):
"""Portable replacement for tf.einsum('bxyc,bzx->bzyc', a, b)."""
expanded_a = tf.expand_dims(a, axis=1)
expanded_b = tf.expand_dims(tf.expand_dims(b, axis=-1), axis=-1)
return tf.reduce_sum(tf.math.multiply(expanded_a,
expanded_b,
name=name),
axis=2)
def FProp(self, theta, inputs, paddings):
"""Apply global spatial pooling to inputs.
Args:
theta: A `.NestedMap` object containing weights' values of this layer and
its children layers.
inputs: The inputs tensor. It is expected to be of shape [batch, time,
frequency, channel]. The time dimension corresponds to the height
dimension as in images and the frequency dimension corresponds to the
width dimension as in images.
paddings: The paddings tensor. It is expected to be of shape [batch,
time]. Defaults to None, which means there no paddings.
Returns:
outputs, out_paddings pair.
- outputs: has shape [batch, 1, 1, channel].
- out_paddings: None or has shape [batch, 1].
"""
p = self.params
assert p.pooling_type in ['MAX', 'AVG'], p.pooling_type
b, t, f = py_utils.GetShape(inputs, ndims=3)
if paddings is not None:
paddings = py_utils.HasShape(paddings, [b, t])
if paddings is not None:
mask = 1.0 - paddings[..., tf.newaxis, tf.newaxis]
else:
mask = tf.ones([b, t, 1, 1], p.dtype)
if p.pooling_type == 'AVG':
global_sum = tf.reduce_sum(inputs * mask, axis=[1, 2], keepdims=True)
f = tf.cast(tf.convert_to_tensor(f), p.dtype)
count = f * tf.reduce_sum(mask, axis=[1, 2], keepdims=True)
out_feature = global_sum / tf.maximum(1.0, count)
elif p.pooling_type == 'MAX':
large_negative = (
tf.ones_like(inputs) * p.dtype.max * tf.constant(-0.7, dtype=p.dtype))
padded_inputs = tf.where_v2(mask > 0.0, inputs, large_negative)
out_feature = tf.reduce_max(padded_inputs, axis=[1, 2], keepdims=True)
if paddings is None:
out_paddings = None
else:
out_paddings = tf.reduce_min(paddings, axis=1, keepdims=True)
out_feature *= 1.0 - out_paddings[..., tf.newaxis, tf.newaxis]
return out_feature, out_paddings
def Proc(record):
"""Parses a serialized tf.Example record."""
outputs = [
('source_id', tf.io.VarLenFeature(tf.int64)),
('source_padding', tf.io.VarLenFeature(tf.float32)),
('target_id', tf.io.VarLenFeature(tf.int64)),
('target_padding', tf.io.VarLenFeature(tf.float32)),
('target_label', tf.io.VarLenFeature(tf.int64)),
('target_weight', tf.io.VarLenFeature(tf.float32)),
]
features = tf.io.parse_single_example(record, dict(outputs))
for k, v in six.iteritems(features):
features[k] = v.values
bucket_key = tf.cast(
tf.maximum(tf.reduce_sum(1.0 - features['source_padding']),
tf.reduce_sum(1.0 - features['target_padding'])),
tf.int32)
return [features[k] for k, _ in outputs], bucket_key
def _InputBatch(self):
length = tf.reduce_prod(self.shape)
counter = summary_utils.StatsCounter('CountingInputGenerator')
new_value = tf.cast(counter.IncBy(length), dtype=tf.int32) - length
new_value = tf.stop_gradient(new_value)
values = new_value + tf.range(length)
shaped_values = tf.reshape(tf.cast(values, dtype=tf.float32),
self.shape)
targets = tf.reduce_sum(shaped_values, axis=0)
return py_utils.NestedMap(src_ids=shaped_values, tgt_ids=targets)
def UnstackFeatures(self, src_inputs, src_paddings):
"""Unstacks src_input and src_paddings based off stack height."""
sh = self.params.stack_height
bs, old_series_length, _, channels = py_utils.GetShape(src_inputs)
unstacked_series_length = old_series_length * sh
src_inputs = tf.reshape(src_inputs,
[bs, unstacked_series_length, -1, channels])
content = 1 - src_paddings
lengths = tf.cast(sh * tf.reduce_sum(content, axis=1), tf.int32)
mask = tf.sequence_mask(lengths, maxlen=unstacked_series_length)
src_paddings = 1 - tf.cast(mask, tf.int32)
return src_inputs, src_paddings
def _AugmentationNetwork(self,
series_length,
inputs,
paddings,
global_seed,
domain_id_index=0):
"""Returns augmented features.
Args:
series_length: Total length of time series.
inputs: Batch of input features of shape (batch_size, time_length,
num_freq, channels).
paddings: Batch of padding vectors of shape (batch_size, time_length).
global_seed: an integer seed tensor for stateless random ops.
domain_id_index: domain id index.
Returns:
Batch of output features of shape (batch_size, time_length, num_freq,
channels) obtained by applying random augmentations to inputs.
"""
p = self.params
dtype = p.dtype
# Unstack the features.
if p.unstack:
inputs, paddings = self.UnstackFeatures(inputs, paddings)
lengths = tf.reduce_sum(1 - paddings, 1)
inputs = self._TimeWarp(inputs,
lengths,
global_seed=global_seed,
dtype=dtype,
domain_id_index=domain_id_index)
inputs = self._TimeMask(inputs,
lengths,
global_seed=global_seed,
noisify=p.use_noise,
gaussian_noise=p.gaussian_noise,
dtype=dtype,
domain_id_index=domain_id_index)
inputs = self._FrequencyMask(inputs,
global_seed=global_seed,
dtype=dtype,
domain_id_index=domain_id_index)
# Restack the features after applying specaugment.
if p.unstack:
inputs = tf.reshape(
inputs,
[tf.shape(inputs)[0], series_length, -1,
tf.shape(inputs)[3]])
return inputs
def _GetExpertDist(self, theta, inputs, *args):
"""Get the task id from inputs tensors."""
# TODO(huangyp): support the more general case when batch size is not 1.
# Input shape can be either [batch, length, dim] or [length, batch, dim]
reshaped_inputs = tf.reshape(inputs, [-1, self.params.cond_dim])
if self.params.nonzeros_mean:
per_example_emb = tf.reduce_sum(reshaped_inputs, 0)
nonzeros = tf.cast(
tf.math.count_nonzero(reshaped_inputs, 0), dtype=tf.float32)
per_example_emb /= (nonzeros + 1e-10)
else:
per_example_emb = tf.reduce_mean(reshaped_inputs, 0)
expert_dist = tf.nn.sigmoid(tf.einsum('i,ij->j', per_example_emb, theta.w))
return expert_dist
def SequenceLength(padding):
"""Computes the length of a sequence based on binary padding.
Args:
padding: A tensor of binary paddings shaped [batch, seqlen].
Returns:
seq_lens, A tensor of shape [batch] containing the non-padded length of each
element of plot_tensor along the batch dimension.
"""
seq_lens = tf.cast(tf.round(tf.reduce_sum(1 - padding, axis=1)), tf.int32)
# Get rid of any extra dimensions.
batch_size = tf.shape(padding)[0]
seq_lens = tf.reshape(seq_lens, [batch_size], name='seq_lens')
return seq_lens
def _global_seed_from_inputs(input_floats):
"""Generates a random seed tensor based on input floats and mode key.
Args:
input_floats: a set of float input tensors that are derived from the input
data (for example, input tokens). The important thing is that these are
usually different for each batch.
Returns:
A tensor of shape=[2] with integer seed tensors derived from the inputs.
"""
timestamp = tf.math.floormod(tf.cast(tf.timestamp(), dtype=tf.int64),
10000000)
input_sum = tf.cast(tf.reduce_sum(tf.math.abs(input_floats)),
dtype=tf.int64)
return tf.stack([timestamp + input_sum, timestamp - input_sum], axis=-1)
def ComputeLoss(self, theta, predicted, input_batch):
diff = predicted - input_batch.tgt_ids
per_example_loss = diff * diff
batch_dim = py_utils.GetShape(per_example_loss)[0]
def replicate_var(name):
return tf.convert_to_tensor([self._private_vars[name]] * batch_dim,
dtype=tf.float32)
metrics = {'loss': (tf.reduce_sum(per_example_loss), batch_dim)}
per_example_tensors = {
'input': input_batch.src_ids,
'loss': per_example_loss,
'diff': diff,
'm': replicate_var('m'),
'b': replicate_var('b'),
}
return metrics, per_example_tensors
def SequenceTrimLastToken(x, x_paddings):
"""Trims the last token off of sequence `x`, and set trimmed elements to 0.
Args:
x: A sequence of tokens of shape [batch_size, x_len_max].
x_paddings: The paddings of `x`.
Returns:
A tuple.
- The new sequence, Tensor of shape [batch_size, x_len_max].
- The new paddings, Tensor of shape [batch_size, x_len_max].
"""
x_len = tf.reduce_sum(1 - x_paddings, 1)
x_len_max = py_utils.GetShape(x)[1]
x_trimmed_len = tf.maximum(x_len - 1, 0)
x_trimmed_paddings = tf.sequence_mask(x_trimmed_len, x_len_max,
x_paddings.dtype)
x_trimmed = x * tf.cast(x_trimmed_paddings, x.dtype)
return x_trimmed, 1 - x_trimmed_paddings
def _ProcessMASSInput(self, source_id, src):
"""Perform MASS input processing."""
# TODO(yuancao): By doing so we assume that right now for monolingual
# eval/dev sets (xx->xx) are in double-column format (since it bypasses
# the Mass op). Ideally we should add a dedicated eval/dev processing
# procedure for unsupervised MT cases, so that single-column eval/devs sets
# are also supported. This should not be handled by any specific ops like
# Mass, but inside the TextPackedInput class.
assert not self.do_eval, 'MASS input can only be used for training.'
_, labels, paddings = self.StringsToIds(tf.reshape(src, [1]),
is_source=True,
key=self._src_tokenizer_key)
weights = 1 - paddings
actual_seq_len = tf.cast(tf.reduce_sum(weights, 1), tf.int32)
src_lang_ids, tgt_lang_ids = self._GetTaskIds(source_id)
mass_out = self.mass_layer.Mask(labels, weights, actual_seq_len)
features = py_utils.NestedMap()
features.src = py_utils.NestedMap()
features.src.ids = mass_out.src.ids
features.src.paddings = paddings
features.src.weights = weights
features.src.task_ids = tf.cast(features.src.weights,
dtype=tf.int32) * src_lang_ids
features.src.ids_indicator = weights
features.tgt = py_utils.NestedMap()
features.tgt.ids = mass_out.tgt.ids
features.tgt.labels = mass_out.tgt.labels
features.tgt.paddings = paddings
features.tgt.weights = mass_out.tgt.weights
features.tgt.task_ids = tf.ones_like(features.src.task_ids,
dtype=tf.int32) * tgt_lang_ids
features.tgt.ids_indicator = weights
if not py_utils.use_tpu():
features.src.strs = src
features.tgt.strs = src
return features.Transform(tf.squeeze)
def FProp(self, theta, input_batch):
"""Embeds source ids and transforms with TransformerStack.
Args:
theta: A `.NestedMap` object containing weights' values of this
layer and its children layers.
input_batch: A `.NestedMap` with fields:
- ids: The inputs tensor. It is expected to be of shape [batch, time].
- paddings: The paddings tensor. Expected shape [batch, time].
- task_ids: If p.task_emb is provided, must contain per-token task
ids of shape [batch, time].
Returns:
A NestedMap containing
- encoded: The encoded features, either a tensor of shape
[time, batch, depth], or a list of tensors if is_transparent is set in
transformer_stack.
- padding: of shape [time, batch]
- segment_id: [time, batch] if packed inputs are supported by the model
(and all layers), or None otherwise.
- embedded_inputs: [time, batch, depth] embedded inputs tokens without
positional encodings.
"""
p = self.params
with tf.name_scope(p.name):
src_segment_id = None
src_segment_pos = None
input_ids = py_utils.with_dependencies([
py_utils.assert_shape_match(tf.shape(input_batch.ids),
tf.shape(input_batch.paddings)),
py_utils.assert_equal(tf.rank(input_batch.ids), 2)
], input_batch.ids)
if (not py_utils.use_tpu()
and tf.flags.FLAGS.transformer_encoder_truncates_inputs):
max_seq_length = tf.cast(
tf.reduce_max(tf.reduce_sum(1.0 - input_batch.paddings,
1)), tf.int32)
paddings = py_utils.with_dependencies([
py_utils.assert_equal(
tf.constant(True, tf.bool),
tf.reduce_all(
input_batch.paddings[:, max_seq_length:] > 0.5))
], input_batch.paddings)
input_ids = input_ids[:, :max_seq_length]
paddings = paddings[:, :max_seq_length]
if p.packed_input:
src_segment_id = input_batch.segment_ids[:, :
max_seq_length]
src_segment_pos = input_batch.segment_pos[:, :
max_seq_length]
else:
paddings = input_batch.paddings
if p.packed_input:
src_segment_id = input_batch.segment_ids
src_segment_pos = input_batch.segment_pos
max_time = tf.shape(input_ids)[1]
# Input token embeddings + positional embeddings
if not p.shared_emb:
input_embs = self.token_emb.EmbLookup(
theta.token_emb, tf.reshape(input_ids, [-1]))
else:
input_embs = self.softmax.EmbLookup(
theta.softmax, tf.reshape(input_ids, [-1]))
input_embs = tf.reshape(input_embs,
[-1, max_time, p.token_emb.embedding_dim])
# [time, batch, dim]
orig_input_embs = tf.transpose(input_embs, [1, 0, 2])
if p.packed_input:
position_embs = self.position_emb.FPropWithPosition(
theta.position_emb, src_segment_pos)
else:
position_embs = self.position_emb.FProp(
theta.position_emb, max_time)
position_embs = tf.reshape(
position_embs, [1, max_time, p.token_emb.embedding_dim])
input_embs += position_embs
if p.task_emb:
input_embs += self.task_emb.EmbLookup(theta.task_emb,
input_batch.task_ids)
if p.model_dim != p.token_emb.embedding_dim:
input_embs = self.emb_proj.FProp(theta.emb_proj, input_embs)
paddings = tf.cast(tf.transpose(paddings), py_utils.FPropDtype(p))
if p.packed_input:
src_segment_id = tf.transpose(src_segment_id)
input_embs = self.input_dropout.FProp(theta.input_dropout,
input_embs)
# [time, batch, dim]
transformer_input = tf.transpose(input_embs, [1, 0, 2])
if not self.do_eval and p.apply_source_mask:
# Augment padding for masked source word positions.
dtype = paddings.dtype
source_mask = tf.where(tf.equal(input_ids, p.source_mask_id),
tf.ones_like(input_ids, dtype=dtype),
tf.zeros_like(input_ids, dtype=dtype))
# Make sure padding is between 0 and 1.
paddings = tf.clip_by_value(paddings + tf.transpose(source_mask),
0.0, 1.0)
encoded, padding, segment_id = self.transformer_stack.FProp(
theta.transformer_stack, transformer_input, paddings,
src_segment_id)
return py_utils.NestedMap(encoded=encoded,
padding=padding,
segment_id=segment_id,
embedded_inputs=orig_input_embs)
def ComputePredictions(self, theta, input_batch):
"""sum(m * x) + b."""
return tf.reduce_sum(theta.m * input_batch.src_ids, axis=1) + theta.b
def _StringsToIdsImpl(self, strs, max_length, append_eos, languages):
"""Takes a tensor of strings and returns id/padding tensors.
This generates `token_ids`, `target_ids`, and `paddings` in the format that
is expected for tokenizers. This performs padding to a fixed length and
appends the end-of-sentence token as appropriate.
Args:
strs: a string Tensor.
max_length: a python integer. The second dimension of the returned arrays.
All sequences are padded or truncated to that length.
append_eos: a python bool. See `BaseTokenizer` for explanation.
languages: A vector of strings with the same length as `strs`.
Returns:
A tuple of 3 tensors:
- token_ids: a tensor of sequences of WPM ids starting with SOS. Sequences
always end with EOS unless the sequence exceeds the maximum length.
Always padded with EOS.
- target_ids: a tensor of sequences of WPM ids not starting with SOS
but ending with EOS. Always padded with EOS.
- paddings: a tensor of floats indicating, at each position, whether
the corresponding position is padded.
"""
p = self.params
if append_eos is None:
append_eos = p.append_eos
batch_size = py_utils.GetShape(strs)[0]
token_ids_ta = tf.TensorArray(tf.int32, batch_size)
target_ids_ta = tf.TensorArray(tf.int32, batch_size)
paddings_ta = tf.TensorArray(tf.float32, batch_size)
def _TokenizeOneSentence(i, strs, token_ids_ta, target_ids_ta, paddings_ta):
"""Tokenizes a single sentence."""
ids, _ = self._wpm_encoder.Encode(strs[i])
if append_eos:
ids = tf.concat([ids, [self.eos_id]], axis=0)
# This truncates after the eos is added, so some sentences might
# not have </s> at the end.
token_ids_ta = token_ids_ta.write(
i,
py_utils.PadOrTrimTo(
tf.concat([[self.sos_id], ids], axis=0), [max_length],
self.eos_id))
target_ids_ta = target_ids_ta.write(
i, py_utils.PadOrTrimTo(ids, [max_length], self.eos_id))
paddings_ta = paddings_ta.write(
i,
py_utils.PadOrTrimTo(
tf.zeros_like(ids, dtype=tf.float32), [max_length], 1.))
return i + 1, strs, token_ids_ta, target_ids_ta, paddings_ta
_, _, token_ids_ta, target_ids_ta, paddings_ta = tf.while_loop(
lambda i, *_: i < batch_size,
_TokenizeOneSentence,
loop_vars=(tf.constant(0, tf.int32), strs, token_ids_ta, target_ids_ta,
paddings_ta),
parallel_iterations=30,
back_prop=False)
token_ids = token_ids_ta.stack()
target_ids = target_ids_ta.stack()
paddings = paddings_ta.stack()
if not p.pad_to_max_length:
maxlen = tf.cast(
tf.round(tf.reduce_max(tf.reduce_sum(1.0 - paddings, axis=1))),
tf.int32)
token_ids = token_ids[:, :maxlen]
target_ids = target_ids[:, :maxlen]
paddings = paddings[:, :maxlen]
return token_ids, target_ids, paddings
def _CreateCanvasAndTargets(self, batch):
# pyformat: disable
"""Create the canvas and targets.
Args:
batch: A `.NestedMap`.
- src: A `.NestedMap`.
- ids: The source ids, ends in <eos>.
- paddings: The source paddings.
- tgt: A `.NestedMap`.
- ids: The target ids, ends in <eos>.
- paddings: The target paddings.
Returns:
A `NestedMap`.
- canvas: The canvas (based off of the `rollin_policy`) of shape
[batch_size, c_dim].
- canvas_paddings: The paddings of `canvas_indices`.
- target_indices: The target indices (i.e., use these indices to
tf.gather_nd the log-probs). Optional, only during training.
- target_weights: The target weights. Optional, only during training.
"""
# pyformat: enable
p = self.params
if not self.do_eval:
# Sample our src and tgt canvas.
src_descriptor = self._SampleCanvasAndTargets(
batch.src.ids, batch.src.paddings)
tgt_descriptor = self._SampleCanvasAndTargets(
batch.tgt.ids, batch.tgt.paddings)
# Offset the src ids (to unshare embeddings between src/tgt). Note, we
# only offset the canvas ids, but we do not offset the vocab ids. This
# will result in unshared embeddings, but shared softmax. This is due to
# GPU/TPU memory limitations, empirically it is known that unsharing
# everything results in better performance.
vocab_size = p.decoder.softmax.num_classes
src_descriptor.canvas = tf.where(
tf.equal(src_descriptor.canvas_paddings, 0),
src_descriptor.canvas + vocab_size, src_descriptor.canvas)
# Offset the tgt indices (need shift according to src length).
batch_size = py_utils.GetShape(batch.src.ids)[0]
# `target_batch` is a [num_targets, batch_size] tensor where each row
# identifies which batch the target belongs to. Note the observation that,
# tf.reduce_sum(target_batch, 1) == 1 \forall rows.
target_batch = tf.cast(
tf.equal(
tf.expand_dims(tf.range(batch_size), 0),
tf.expand_dims(tgt_descriptor.target_indices[:, 0], 1)),
tf.int32)
src_lens = tf.cast(
tf.reduce_sum(1 - src_descriptor.canvas_paddings, 1), tf.int32)
# `tgt_offset` is shape [num_targets] where each entry corresponds to the
# offset needed for that target (due to the source length).
tgt_offset = tf.matmul(target_batch, tf.expand_dims(src_lens, 1))
# We shift the tgt slot without touching the batch or vocab.
tgt_descriptor.target_indices += tf.concat([
tf.zeros_like(tgt_offset), tgt_offset,
tf.zeros_like(tgt_offset)
], 1)
# The canvas is simply the sequence-level concat of the src and tgt.
canvas, canvas_paddings = insertion.SequenceConcat(
src_descriptor.canvas, src_descriptor.canvas_paddings,
tgt_descriptor.canvas, tgt_descriptor.canvas_paddings)
target_indices = tf.concat(
[src_descriptor.target_indices, tgt_descriptor.target_indices],
0)
target_weights = tf.concat(
[src_descriptor.target_weights, tgt_descriptor.target_weights],
0)
return py_utils.NestedMap(canvas=canvas,
canvas_paddings=canvas_paddings,
target_indices=target_indices,
target_weights=target_weights)
def FProp(self,
theta,
x,
x_paddings=None,
eos_id=1,
force_sample_last_token=True):
"""Applies SymbolInsertionLayer.
We take in a `x`, which represents the groundtruth sequence (i.e., English
sequence). We return a sampled rollin (observed) canvas (i.e., random subset
of the English sequence), as well as the target (indices) for an
insertion-based model (i.e., the targets given the random observed subset).
Args:
theta: Ignored, this can be None.
x: The symbol ids of shape `[batch_size, time_dim]`.
x_paddings: The paddings (1 or 0) of shape `[batch_size, time_dim]` where
0 is valid and 1 is invalid.
eos_id: The <eos> token id to represent end-of-slot.
force_sample_last_token: Set True to force sample the last token of `x`.
Returns:
A `NestedMap`.
- canvas: The canvas (based off of the `rollin_policy`) of shape
[batch_size, c_dim]. Note that, `c_dim` <= `time_dim` but need not be
equal.
- canvas_indices: The canvas indices (into `x`).
- canvas_paddings: The paddings of `canvas_indices`.
- target_indices: The target indices of shape [num_targets, 3].
`num_targets` is the number of total targets in the entire batch.
[:, 0] captures the batch, [:, 1] captures the slot, and [:, 2]
captures the token. Each row [batch, slot, vocab] represents the
indices of the target -- i.e., the batch, slot and vocab combination
of the target. Typical usage of these indices is to tf.gather_nd
the log-probs (from the softmax layer).
- target_weights: The target weights.
Raises:
ValueError: If invalid params.
"""
p = self.params
batch_size = py_utils.GetShape(x)[0]
time_dim = py_utils.GetShape(x)[1]
if x_paddings is None:
x_paddings = tf.zeros([batch_size, time_dim], tf.float32)
oracle_policy = p.oracle_policy
rollin_policy = (oracle_policy
if p.rollin_policy == 'oracle' else p.rollin_policy)
if rollin_policy != 'uniform':
raise ValueError('Unknown or unsupported rollin policy: %s' %
rollin_policy)
if oracle_policy != 'uniform':
raise ValueError('Unknown or unsupported oracle policy: %s' %
oracle_policy)
x_len = tf.cast(tf.round(tf.reduce_sum(1 - x_paddings, 1)), tf.int32)
# Compute the desired length per example in the batch.
ratio = tf.random.uniform([batch_size], 0.0, 1.0, seed=p.random_seed)
if force_sample_last_token:
c_len = tf.minimum(
tf.cast(ratio * tf.cast(x_len, tf.float32), tf.int32),
x_len - 1) + 1
else:
c_len = tf.minimum(
tf.cast(ratio * tf.cast(x_len + 1, tf.float32), tf.int32),
x_len)
# Compute the maximum length across the batch.
c_len_max = tf.reduce_max(c_len)
# Grab subset of random valid indices per example.
z_logits = tf.cast(
tf.expand_dims(tf.range(time_dim), 0) >= tf.expand_dims(x_len, 1),
tf.float32) * -1e9
if force_sample_last_token:
# Force sample the last token -- i.e., as indexed by `x_len - 1`. We can
# accomplish this by add +LARGE_NUMBER to the logits.
z_logits += tf.cast(
tf.equal(tf.expand_dims(tf.range(time_dim), 0),
tf.expand_dims(x_len - 1, 1)), tf.float32) * 1e9
# Gumbel-max trick to sample (we only sample valid positions per sample in
# the batch).
z = -tf.math.log(-tf.math.log(
tf.random.uniform([batch_size, time_dim], seed=p.random_seed)))
unused_c_values, c_indices = tf.nn.top_k(z_logits + z, time_dim)
# Trim everything > c_len_max.
c_indices = c_indices[:, :c_len_max]
# Invalidate any indices >= c_len, we use the last index as the default
# invalid index.
c_indices = tf.where(
tf.expand_dims(tf.range(c_len_max), 0) < tf.expand_dims(c_len, 1),
c_indices, tf.fill(py_utils.GetShape(c_indices), time_dim - 1))
# Materialize the canvas.
c_indices = tf.sort(c_indices)
c = tf.gather_nd(
x,
tf.stack([
tf.reshape(
tf.tile(tf.expand_dims(tf.range(batch_size), 1),
[1, c_len_max]), [-1]),
tf.reshape(c_indices, [-1])
], 1))
c = tf.reshape(c, [batch_size, c_len_max])
# Compute the paddings.
c_paddings = 1 - tf.sequence_mask(
c_len, c_len_max, dtype=x_paddings.dtype)
c *= tf.cast(1 - c_paddings, tf.int32)
indices = tf.concat([
tf.reshape(
tf.tile(tf.expand_dims(tf.range(batch_size), 1),
[1, c_len_max]), [batch_size * c_len_max, 1]),
tf.reshape(c_indices, [batch_size * c_len_max, 1])
], 1)
x_token_is_observed = tf.scatter_nd(
indices, tf.ones([batch_size * c_len_max], tf.int32),
py_utils.GetShape(x))
# `x_segments` captures which slot each `x` belongs to (both observed and
# tokens that need to be observed).
x_segments = tf.cumsum(x_token_is_observed, 1, exclusive=True)
x_token_is_observed = tf.cast(x_token_is_observed, tf.bool)
prev_x_token_is_observed = tf.pad(x_token_is_observed[:, :-1],
[[0, 0], [1, 0]],
constant_values=True)
x_token_is_observed = tf.reshape(x_token_is_observed, [-1])
prev_x_token_is_observed = tf.reshape(prev_x_token_is_observed, [-1])
x_is_valid = tf.cast(1 - x_paddings, tf.bool)
x_is_valid = tf.reshape(x_is_valid, [-1])
# Remap all the observed to <eos>, note some of these need a zero weight
# (or else there would be <eos> and valid token in the same slot).
target_indices = tf.cast(tf.reshape(x, [-1, 1]), tf.int32)
target_indices = tf.where(
x_token_is_observed,
tf.fill(py_utils.GetShape(target_indices), eos_id), target_indices)
# TODO(williamchan): We give uniform 1.0 weight, however, math suggests
# we may want to weigh this term by the original sequence length.
target_weights = tf.ones_like(target_indices, tf.float32)
# We need to set all the weights for <eos> which actually have valid tokens
# in the slot to zero.
target_weights = tf.where(
x_token_is_observed & ~prev_x_token_is_observed,
tf.zeros_like(target_weights), target_weights)
# TODO(williamchan): Consider dropping the entries w/ weight zero.
# Add the batch and slot indices.
target_indices = tf.concat([
tf.reshape(
tf.tile(tf.expand_dims(tf.range(batch_size), 1),
[1, time_dim]), [batch_size * time_dim, 1]),
tf.reshape(x_segments, [-1, 1]), target_indices
], 1)
# Select only the valid indices. The selected valid ones include slots w/
# <eos>.
target_indices = target_indices[x_is_valid]
target_weights = target_weights[x_is_valid]
return py_utils.NestedMap(canvas=c,
canvas_indices=c_indices,
canvas_paddings=c_paddings,
target_indices=target_indices,
target_weights=target_weights)
def GlobalBatchSize(self):
"""Returns the total number of examples in the current batch."""
# The number of examples is indicated by the segment_ids of the target.
num_segments = tf.math.reduce_max(self._batch.tgt.segment_ids, axis=1)
return tf.reduce_sum(tf.cast(num_segments, dtype=tf.int32))
def BeamSearchDecode(self,
theta,
encoder_outputs,
num_hyps_per_beam_override=0,
init_beam_search_state=None,
pre_beam_search_step_callback=None,
post_beam_search_step_callback=None,
max_steps=None):
"""Performs beam-search based decoding.
Args:
theta: A NestedMap object containing weights' values of the decoder layer
and its children layers.
encoder_outputs: A NestedMap containing encoder outputs to be passed to
the callbacks. Mostly opaque to BeamSearchHelper, except that it should
contain either a 'seq_lengths' field of shape [source_batch_size] or
a 'paddings' field of shape [source_max_lengths, source_batch_size].
num_hyps_per_beam_override: If set to a value <= 0, this parameter is
ignored. If set to a value > 0, then this value will be used to override
`p.num_hyps_per_beam`.
init_beam_search_state: The `InitBeamSearchState` callback. Please refer
to the class header comments for more details.
pre_beam_search_step_callback: The `PreBeamSearchStepCallback` callback.
Please refer to the class header comments for more details.
post_beam_search_step_callback: The `PostBeamSearchStepCallback` callback.
Please refer to the class header comments for more details.
max_steps: maximum beam search steps. If None, use
self.params.target_seq_len.
Returns:
A `BeamSearchDecodeOutput`.
"""
p = self.params
num_hyps_per_beam = p.num_hyps_per_beam
if num_hyps_per_beam_override > 0:
num_hyps_per_beam = num_hyps_per_beam_override
if max_steps is None:
max_steps = p.target_seq_len
initial_results, other_states = init_beam_search_state(
theta, encoder_outputs, num_hyps_per_beam)
num_hyps = tf.shape(initial_results.log_probs)[0]
num_beams = num_hyps // num_hyps_per_beam
if 'step_ids' in initial_results:
# [num_hyps, 1]
step_ids = tf.ensure_shape(initial_results.step_ids, [None, 1])
else:
step_ids = tf.fill([num_hyps, 1],
tf.constant(p.target_sos_id, dtype=tf.int32))
min_score = -1e36
best_scores = (tf.zeros(shape=[num_beams], dtype=p.dtype) + min_score)
cumulative_scores = tf.zeros(shape=[num_hyps], dtype=p.dtype)
in_scores = tf.zeros([max_steps, num_hyps], dtype=p.dtype)
in_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.int32)
in_prev_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.int32)
in_done_hyps = tf.zeros([max_steps, num_hyps], dtype=tf.string)
bs_atten_probs = tf.zeros(
[max_steps, num_hyps,
tf.shape(initial_results.atten_probs)[1]],
dtype=p.dtype)
cur_step = tf.constant(0, dtype=tf.int32)
all_done = tf.constant(False, dtype=tf.bool)
core_bs_states = (best_scores, cumulative_scores, in_scores, in_hyps,
in_prev_hyps, in_done_hyps, bs_atten_probs)
def LoopContinue(cur_step, all_done, unused_step_ids,
unused_core_bs_states, unused_other_states_list):
return tf.math.logical_and(cur_step < max_steps,
tf.math.logical_not(all_done))
def LoopBody(cur_step, unused_all_done, step_ids, core_bs_states,
other_states_list):
(cur_step, all_done, new_step_ids, new_bs_states,
new_other_states) = self._BeamSearchStep(
theta, encoder_outputs, cur_step, step_ids, core_bs_states,
other_states.Pack(other_states_list), num_hyps_per_beam,
pre_beam_search_step_callback, post_beam_search_step_callback)
return (cur_step, all_done, new_step_ids, new_bs_states,
new_other_states.Flatten())
flat_other_states = other_states.Flatten()
_, _, _, final_bs_states, flat_final_other_states = tf.while_loop(
LoopContinue,
LoopBody,
loop_vars=(cur_step, all_done, step_ids, core_bs_states,
flat_other_states),
parallel_iterations=10,
back_prop=False,
swap_memory=False,
shape_invariants=(tf.TensorShape(cur_step.get_shape()),
tf.TensorShape(all_done.get_shape()),
tf.TensorShape(step_ids.get_shape()),
_GetShapes(core_bs_states),
_GetShapes(flat_other_states, none_shapes=True)))
# [target_seq_len, num_beams * num_hyps_per_beam].
final_done_hyps = final_bs_states[5]
final_other_states = other_states.Pack(flat_final_other_states)
# Assume that `paddings` has shape [source_max_lengths, source_batch_size]
# by default, and compute `encoded_seq_lengths` accordingly. This can be
# overridden by directly passing `seq_lengths` in the `encoder_outputs`
# NestedMap.
encoded_seq_lengths = getattr(encoder_outputs, 'seq_lengths', None)
if encoded_seq_lengths is None:
source_paddings = encoder_outputs.padding
if isinstance(source_paddings, py_utils.NestedMap):
encoded_seq_lengths = tf.cast(
tf.round(
tf.reduce_sum(
1.0 - tf.transpose(source_paddings.Flatten()[0]),
1)), tf.int32)
else:
encoded_seq_lengths = tf.cast(
tf.round(
tf.reduce_sum(1.0 - tf.transpose(source_paddings), 1)),
tf.int32)
# [num_beams, num_hyps_per_beam].
topk_hyps = ops.top_k_terminated_hyps(
final_done_hyps,
encoded_seq_lengths,
k=num_hyps_per_beam,
num_hyps_per_beam=num_hyps_per_beam,
length_normalization=p.length_normalization,
coverage_penalty=p.coverage_penalty,
target_seq_length_ratio=p.target_seq_length_ratio,
eoc_id=p.target_eoc_id,
merge_paths=p.merge_paths)
# [num_beams * num_hyps_per_beam, ...].
max_seq_length = 0 if isinstance(max_steps, tf.Tensor) else max_steps
topk_ids, topk_lens, topk_scores = ops.unpack_hyp(
tf.reshape(topk_hyps, [-1]), max_seq_length=max_seq_length)
# [num_beams, num_hyps_per_beam].
topk_scores = tf.reshape(topk_scores, tf.shape(topk_hyps))
return BeamSearchDecodeOutput(final_done_hyps, topk_hyps, topk_ids,
topk_lens, topk_scores, None,
final_other_states)
def _BeamSearchDecodeIds(self,
theta,
encoder_outputs,
num_hyps_per_beam,
init_beam_search_state=None,
pre_beam_search_step_callback=None,
post_beam_search_step_callback=None,
max_steps=None):
"""Performs beam-search based decoding.
Args:
theta: A NestedMap object containing weights' values of the decoder layer
and its children layers.
encoder_outputs: A NestedMap computed by encoder.
num_hyps_per_beam: Number of hyps per beam.
init_beam_search_state: The InitBeamSearchState callback. Please refer to
the class header comments for more details.
pre_beam_search_step_callback: The PreBeamSearchStepCallback callback.
Please refer to the class header comments for more details.
post_beam_search_step_callback: The PostBeamSearchStepCallback callback.
Please refer to the class header comments for more details.
max_steps: maximum beam search steps. If None, use
self.params.target_seq_len.
Returns:
hyps: A tensor of shape [time, b * k] with ids of the token selected.
prev_hyps: A tensor of shape [time, b * k] with index to the previous hyps
which was selected.
done_hyps: A boolean tensor of shape [time, b * k] where value
indicates if hyps was terminated.
scores: A tensor of shape [time, b * k] with scores of the token
selected.
atten_probs: A tensor of shape [time, b * k, seq_len] which contain the
attention probabilities over the source words against word in the
previous hyps.
eos_scores: A tensor of shape [time, b * k] with scores of the eos token
selected.
eos_atten_probs: A tensor of shape [time, b * k, seq_len] which contain
the attention probabilities over the source words against word in the
previous hyps.
source_seq_lengths: A tensor of shape [time] containing the source
seq_lengths.
flat_final_other_states: A array of tensors that are part of other states.
"""
p = self.params
source_paddings = encoder_outputs.padding
initial_results, other_states = init_beam_search_state(
theta, encoder_outputs, num_hyps_per_beam)
num_hyps = tf.shape(initial_results.log_probs)[0]
num_beams = num_hyps // num_hyps_per_beam
# We cache the NestedMap as member variable so that we can use it to
# pack the final outputs. Tpu rewrite methods forces us to strictly pass
# in Tensors, and output Tensors
self._other_states = other_states
step_ids = tf.fill([num_hyps, 1],
tf.constant(p.target_sos_id, dtype=tf.int32))
min_score = -1e36
fprop_dtype = py_utils.FPropDtype(p)
best_scores = (tf.zeros(shape=[num_beams], dtype=fprop_dtype) +
min_score)
cumulative_scores = tf.zeros(shape=[num_hyps], dtype=fprop_dtype)
histories = tf.zeros(shape=[num_hyps], dtype=tf.int32)
in_scores = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
in_hyps = tf.TensorArray(dtype=tf.int32, size=max_steps)
in_prev_hyps = tf.TensorArray(dtype=tf.int32, size=max_steps)
in_done_hyps = tf.TensorArray(dtype=tf.int32, size=max_steps)
in_atten_probs = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
in_eos_scores = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
in_eos_atten_probs = tf.TensorArray(dtype=fprop_dtype, size=max_steps)
cur_step = tf.constant(0, dtype=tf.int32)
all_done = tf.constant(False, dtype=tf.bool)
# States for beam search that are inputs into Beam search step.
accum_bs_states = [best_scores, cumulative_scores, histories]
# States that are not accumulators.
non_accum_bs_states = [
in_scores,
in_hyps,
in_prev_hyps,
in_done_hyps,
in_atten_probs,
in_eos_scores,
in_eos_atten_probs,
]
core_bs_states = tuple(accum_bs_states + non_accum_bs_states)
flat_other_states = other_states.Flatten()
# If there is an optimized implementation for short sequence, LoopBodyShort
# will run first for short_seq_limit steps (after which the
# LoopBodyShort does not have performance benefit). Then LoopBodyLong (the
# default implementation) is used to continue the rest of the steps. For
# decoders which do not have the short sequence specific implementation,
# only the LoopBodyLong (the default implementation) will run.
if p.short_seq_limit > 0:
def LoopContinueShort(cur_step, all_done, unused_step_ids,
unused_core_bs_states,
unused_other_states_list):
"""Use short_seq optimization when cur_step is smaller than limit."""
return tf.math.logical_and(cur_step < p.short_seq_limit,
tf.math.logical_not(all_done))
def LoopBodyShort(cur_step, unused_all_done, step_ids,
core_bs_states, other_states_list):
"""Loop body of short_seq optimization.
Instead of doing computation for the entire padded sequence, while loop
with early exit is used within each _BeamSearchStep to do computation
for only the actual sequence (seq_length <= cur_step).
use_short_seq_opt is used as the flag to pass this information down to
the decoder implementation.
Args:
cur_step: A scalar int tensor, the current time step, 0-based.
unused_all_done: A tf.bool, indicating whether the decoding finishes.
step_ids: An int32 tensor of shape [num_hyps, 1]. The input ids to the
current search step.
core_bs_states: A tuple of core beam search states.
other_states_list: A flattened NestedMap of other beam search states.
Returns:
The updated input tuple, with the same shape.
"""
(cur_step, all_done, new_step_ids, new_bs_states,
new_other_states) = self._BeamSearchStep(
theta,
encoder_outputs,
cur_step,
step_ids,
core_bs_states,
other_states.Pack(other_states_list),
num_hyps_per_beam,
pre_beam_search_step_callback,
post_beam_search_step_callback,
use_short_seq_opt=True)
return (cur_step, all_done, new_step_ids, new_bs_states,
new_other_states.Flatten())
(cur_step, all_done, step_ids, core_bs_states,
flat_other_states) = tf.while_loop(
LoopContinueShort,
LoopBodyShort,
loop_vars=(cur_step, all_done, step_ids, core_bs_states,
flat_other_states),
parallel_iterations=10,
back_prop=False,
swap_memory=False,
shape_invariants=(
tf.TensorShape(cur_step.get_shape()),
tf.TensorShape(all_done.get_shape()),
tf.TensorShape(step_ids.get_shape()),
tuple(
list(_GetShapes(accum_bs_states)) +
list(_GetShapes(non_accum_bs_states,
none_shapes=True))),
_GetShapes(flat_other_states, none_shapes=True)),
maximum_iterations=max_steps)
def LoopContinueLong(cur_step, all_done, unused_step_ids,
unused_core_bs_states, unused_other_states_list):
"""Continue default implementation until decoding finishes."""
return tf.math.logical_and(cur_step < max_steps,
tf.math.logical_not(all_done))
def LoopBodyLong(cur_step, unused_all_done, step_ids, core_bs_states,
other_states_list):
"""Loop body of default long_seq implementation."""
(cur_step, all_done, new_step_ids, new_bs_states,
new_other_states) = self._BeamSearchStep(
theta,
encoder_outputs,
cur_step,
step_ids,
core_bs_states,
other_states.Pack(other_states_list),
num_hyps_per_beam,
pre_beam_search_step_callback,
post_beam_search_step_callback,
use_short_seq_opt=False)
return (cur_step, all_done, new_step_ids, new_bs_states,
new_other_states.Flatten())
_, _, _, final_bs_states, flat_final_other_states = tf.while_loop(
LoopContinueLong,
LoopBodyLong,
loop_vars=(cur_step, all_done, step_ids, core_bs_states,
flat_other_states),
parallel_iterations=10,
back_prop=False,
swap_memory=False,
shape_invariants=(
tf.TensorShape(cur_step.get_shape()),
tf.TensorShape(all_done.get_shape()),
tf.TensorShape(step_ids.get_shape()),
tuple(
list(_GetShapes(accum_bs_states)) +
list(_GetShapes(non_accum_bs_states, none_shapes=True))),
_GetShapes(flat_other_states, none_shapes=False)),
maximum_iterations=max_steps)
if isinstance(source_paddings, py_utils.NestedMap):
source_seq_lengths = tf.cast(tf.round(
tf.reduce_sum(1.0 - tf.transpose(source_paddings.Flatten()[0]),
1)),
dtype=tf.int32)
else:
source_seq_lengths = tf.cast(tf.round(
tf.reduce_sum(1.0 - tf.transpose(source_paddings), 1)),
dtype=tf.int32)
# Concatenate all outputs on axis=0.
scores = final_bs_states[3].stack()
hyps = final_bs_states[4].stack()
prev_hyps = final_bs_states[5].stack()
done_hyps = tf.cast(final_bs_states[6].stack(), tf.bool)
atten_probs = final_bs_states[7].stack()
eos_scores = final_bs_states[8].stack()
eos_atten_probs = final_bs_states[9].stack()
rets = (hyps, prev_hyps, done_hyps, scores, atten_probs, eos_scores,
eos_atten_probs, source_seq_lengths)
# TODO(rohananil): Only send a single R1 tensor to host instead of 3 after
# b/111131551 is resolved.
# Canonical shapes for tensors of various. ranks
r_shapes = [
py_utils.GetShape(source_seq_lengths),
py_utils.GetShape(hyps),
py_utils.GetShape(atten_probs)
]
# Reshape all tensors to [-1] to avoid cost of copy due to padding.
rets_r1 = [tf.reshape(r, [-1]) for r in rets]
return tuple(r_shapes) + tuple(rets_r1) + tuple(
flat_final_other_states)
def _GetMask(self,
batch_size,
choose_range,
mask_size,
global_seed,
max_length=None,
masks_per_frame=0.0,
multiplicity=1,
dtype=tf.float32,
max_ratio=1.0):
"""Returns fixed size multi-masks starting from random positions.
A multi-mask is a mask obtained by applying multiple masks.
This function when max_length is given:
1) Sample random mask lengths less than max_length with shape
(batch_size, multiplicity).
2) Truncate lengths to a max of (choose_range * max_ratio),
so that each mask is fully contained within the corresponding sequence.
3) Random sample start points of shape (batch_size, multiplicity)
with in (choose_range - lengths).
4) For each batch, multiple masks (whose number is given by the
multiplicity) are constructed.
5) Return a mask of shape (batch_size, mask_size) where masks are
obtained by composing the masks constructed in step 4).
If masks_per_frame > 0, the number is given by
min(masks_per_frame * choose_range, multiplicity).
If not, all the masks are composed. The masked regions are set to zero.
This function when max_length is not given:
1) Sample random mask lengths less than (choose_range * max_ratio)
with shape (batch_size, multiplicity).
2) Proceed to steps 3), 4) and 5) of the above.
Args:
batch_size: Batch size. Integer number.
choose_range: Range within which the masked entries must lie. Tensor of
shape (batch_size,).
mask_size: Size of the mask. Integer number.
global_seed: an integer seed tensor for stateless random ops.
max_length: Maximum number of allowed consecutive masked entries. Integer
number or None.
masks_per_frame: Number of masks per frame. Float number. If > 0, the
multiplicity of the mask is set to be masks_per_frame * choose_range.
multiplicity: Maximum number of total masks. Integer number.
dtype: Data type.
max_ratio: Maximum portion of the entire range allowed to be masked. Float
number.
Returns:
mask: a fixed size multi-mask starting from a random position with shape
(batch_size, mask_size).
"""
p = self.params
# Non-empty random seed values are only used for testing or when using
# stateless random ops. seed_1 and seed_2 are set separately to avoid
# correlation of mask size and mask position.
if p.use_input_dependent_random_seed:
seed_1 = global_seed + 1
seed_2 = global_seed + 2
elif p.random_seed:
seed_1 = p.random_seed + 1
seed_2 = 2 * p.random_seed
else:
seed_1 = p.random_seed
seed_2 = p.random_seed
# Sample lengths for multiple masks.
if max_length and max_length > 0:
max_length = tf.broadcast_to(tf.cast(max_length, dtype),
(batch_size, ))
else:
max_length = tf.cast(choose_range, dtype=dtype) * max_ratio
random_uniform = _random_uniform_op(p.use_input_dependent_random_seed)
masked_portion = random_uniform(shape=(batch_size, multiplicity),
minval=0.0,
maxval=1.0,
dtype=dtype,
seed=seed_1)
masked_frame_size = self.EinsumBBmBm(max_length, masked_portion)
masked_frame_size = tf.cast(masked_frame_size, dtype=tf.int32)
# Make sure the sampled length was smaller than max_ratio * length_bound.
# Note that sampling in this way was biased
# (shorter sequence may over-masked.)
choose_range = tf.expand_dims(choose_range, -1)
choose_range = tf.tile(choose_range, [1, multiplicity])
length_bound = tf.cast(choose_range, dtype=dtype)
length_bound = tf.cast(max_ratio * length_bound, dtype=tf.int32)
length = tf.minimum(masked_frame_size, tf.maximum(length_bound, 1))
# Choose starting point.
random_start = random_uniform(shape=(batch_size, multiplicity),
maxval=1.0,
seed=seed_2)
start_with_in_valid_range = random_start * tf.cast(
(choose_range - length + 1), dtype=dtype)
start = tf.cast(start_with_in_valid_range, tf.int32)
end = start + length - 1
# Shift starting and end point by small value.
delta = tf.constant(0.1)
start = tf.expand_dims(tf.cast(start, dtype) - delta, -1)
start = tf.tile(start, [1, 1, mask_size])
end = tf.expand_dims(tf.cast(end, dtype) + delta, -1)
end = tf.tile(end, [1, 1, mask_size])
# Construct pre-mask of shape (batch_size, multiplicity, mask_size).
diagonal = tf.expand_dims(
tf.expand_dims(tf.cast(tf.range(mask_size), dtype=dtype), 0), 0)
diagonal = tf.tile(diagonal, [batch_size, multiplicity, 1])
pre_mask = tf.cast(tf.math.logical_and(diagonal < end,
diagonal > start),
dtype=dtype)
# Sum masks with appropriate multiplicity.
if masks_per_frame > 0:
multiplicity_weights = tf.tile(
tf.expand_dims(tf.range(multiplicity, dtype=dtype), 0),
[batch_size, 1])
multiplicity_tensor = masks_per_frame * tf.cast(choose_range,
dtype=dtype)
multiplicity_weights = tf.cast(
multiplicity_weights < multiplicity_tensor, dtype=dtype)
pre_mask = self.EinsumBmtBmBt(pre_mask, multiplicity_weights)
else:
pre_mask = tf.reduce_sum(pre_mask, 1)
mask = tf.cast(1.0 - tf.cast(pre_mask > 0, dtype=dtype), dtype=dtype)
if p.fprop_dtype is not None and p.fprop_dtype != p.dtype:
mask = tf.cast(mask, p.fprop_dtype)
return mask
def FProp(self, theta, input_batch):
"""Embeds source ids and transforms with TransformerStack.
Args:
theta: A `.NestedMap` object containing weights' values of this layer and
its children layers.
input_batch: A `.NestedMap` object containing: ids - The inputs tensor of
shape [batch, time]. paddings - The ids' paddings of shape [batch,
time].
Returns:
A '.NestedMap' object containing:
encoded - The encoded features of shape [time, batch, dim] or [batch,
time, dim], depending p.output_data_format.
padding - The encoded features' padding of shape [time, batch] or
[batch, time].
segment_id - The segmentation of packed inputs of shape [time, batch] or
[batch, time] if it is supported by the model, or None otherwise.
embedded_inputs - The embedded inputs tokens without positional
encodings of shape [time, batch, dim] or [batch, time, dim].
"""
p = self.params
with tf.name_scope(p.name):
# [batch, time]
input_ids = input_batch.ids
# [batch, time]
paddings = input_batch.paddings
# [batch, time]
segment_ids = input_batch.segment_ids if p.packed_input else None
batch = py_utils.GetShape(input_ids)[0]
time = py_utils.GetShape(input_ids)[1]
# Embedding layer.
# [batch, time, dim]
if not p.shared_emb:
input_embs = self.token_emb.EmbLookup(theta.token_emb,
input_ids)
else:
input_embs = self.softmax.EmbLookup(theta.softmax, input_ids)
orig_input_embs = input_embs
# [1, time, dim]
if p.packed_input:
positions = input_batch.segment_pos
position_embs = tf.expand_dims(
self.position_emb.FPropWithPosition(
theta.position_emb, positions), 0)
else:
position_embs = tf.expand_dims(
self.position_emb.FProp(theta.position_emb, time), 0)
# [batch, time, dim]
input_embs += tf.cast(position_embs, tf.bfloat16)
if p.input_dropout_tpl.fprop_dtype:
input_embs = tf.cast(input_embs,
p.input_dropout_tpl.fprop_dtype)
paddings = tf.cast(paddings, p.input_dropout_tpl.fprop_dtype)
input_embs = self.input_dropout.FProp(theta.input_dropout,
input_embs)
# [batch, time, dim]
transformer_input = input_embs
# Explicitly set the input shape of Transformer layers, to avoid
# unknown shape error occurred to tf.einsum on nonTPU devices.
transformer_input = tf.reshape(transformer_input,
[batch, time, p.model_dim])
# Compute self-attention segment mask once.
if p.packed_input:
segment_mask = batch_major_attention.SegmentMask(
segment_ids, segment_ids, dtype=transformer_input.dtype)
else:
segment_mask = tf.zeros([batch, 1, time, time])
encoded, padding = self.transformer_stack.FProp(
theta.transformer_stack, transformer_input, paddings,
segment_mask)
if p.final_layer_norm:
encoded = self.final_ln.FProp(theta.final_ln, encoded)
seq_lengths = tf.cast(tf.reduce_sum(1. - padding, axis=1),
tf.int32)
if p.output_data_format == 'TBC':
encoded = tf.transpose(encoded,
[1, 0, 2]) # [time, batch, dim]
padding = tf.transpose(padding) # [time, batch]
segment_ids = tf.transpose(
segment_ids) if p.packed_input else None
orig_input_embs = tf.transpose(orig_input_embs, [1, 0, 2])
return py_utils.NestedMap(
encoded=encoded,
padding=padding,
seq_lengths=seq_lengths, # used by beam_search_helper.
segment_id=segment_ids,
embedded_inputs=orig_input_embs)
