def _repeat_molecules(self, molecules: tf.Tensor,
char_seq_length: tf.Tensor,
molecule_seq_length: tf.Tensor) -> tf.Tensor:
"""Repeats molecules to make them the same length as the char sequence."""
del molecule_seq_length # Used for contract only.
rate = self.config.downsampling_rate
molecules_without_extra_cls = molecules[:, 1:, :]
# `repeated`: [batch_size, almost_char_seq_len, molecule_hidden_size]
repeated = tf.repeat(molecules_without_extra_cls,
repeats=rate,
axis=-2)
# So far, we've repeated the elements sufficient for any `char_seq_length`
# that's a multiple of `downsampling_rate`. Now we account for the last
# n elements (n < `downsampling_rate`), i.e. the remainder of floor
# division. We do this by repeating the last molecule a few extra times.
last_molecule = molecules[:, -1:, :]
remainder_length = tf.floormod(char_seq_length, rate)
remainder_repeated = tf.repeat(
last_molecule,
# +1 molecule to compensate for truncation.
repeats=remainder_length + rate,
axis=-2)
# `repeated`: [batch_size, char_seq_len, molecule_hidden_size]
return tf.concat([repeated, remainder_repeated], axis=-2)
def _aspect_preserving_width_resize(image, width=512):
# If training on ADE20k
height_i = tf.shape(image)[0]
new_height = height_i - tf.floormod(height_i, 16)
return tf.image.resize_image_with_crop_or_pad(
image, new_height, width)
def _aspect_preserving_width_resize(image, width=512):
height_i = tf.shape(image)[0]
# width_i = tf.shape(image)[1]
# ratio = tf.to_float(width_i) / tf.to_float(height_i)
# new_height = tf.to_int32(tf.to_float(height_i) / ratio)
new_height = height_i - tf.floormod(height_i, 16)
return tf.image.resize_image_with_crop_or_pad(
image, new_height, width)
def _hsvnoise(rgb): hsv = tf.image.rgb_to_hsv(rgb / 255.0) # Needs [0 1] input. h, sv = hsv[..., :1], hsv[..., 1:] h = tf.floormod(1. + h + rnd(*h_add), 1.) # color cycle. pow_, mul, add = 2.0**rnd2(*sv_pow), 2.0**rnd2(*sv_mul), rnd2(*sv_add) sv = sv**pow_ * mul + add hsv = tf.clip_by_value(tf.concat([h, sv], axis=-1), 0, 1) return tf.image.hsv_to_rgb(hsv) * 255.0
def _map_fn(index):
x = tf.floordiv(index, 2)
y = tf.floormod(index, 2)
label = tf.cast(index + 1, tf.float32)
label = tf.reshape(label, [1])
target_dense = tf.stack([x + y, x + y + 1])
return ({KEY_NAME: dense_to_sparse(target_dense, tf.int64)}, label)
def define_predictions(self, features, outputs):
"""Define model predictions."""
predictions = {
"example_id": features["example_id"],
"service_id": features["service_id"],
"is_real_example": features["is_real_example"],
}
# Scores are output for each intent.
# Note that the intent indices are shifted by 1 to account for NONE intent.
predictions["intent_status"] = tf.argmax(
outputs["logit_intent_status"], axis=-1)
# Scores are output for each requested slot.
predictions["req_slot_status"] = tf.sigmoid(
outputs["logit_req_slot_status"])
# For categorical slots, the status of each slot and the predicted value are
# output.
predictions["cat_slot_status"] = tf.argmax(
outputs["logit_cat_slot_status"], axis=-1)
predictions["cat_slot_value"] = tf.argmax(
outputs["logit_cat_slot_value"], axis=-1)
# For non-categorical slots, the status of each slot and the indices for
# spans are output.
predictions["noncat_slot_status"] = tf.argmax(
outputs["logit_noncat_slot_status"], axis=-1)
start_scores = tf.nn.softmax(outputs["logit_noncat_slot_start"],
axis=-1)
end_scores = tf.nn.softmax(outputs["logit_noncat_slot_end"], axis=-1)
_, max_num_slots, max_num_tokens = end_scores.get_shape().as_list()
batch_size = tf.shape(end_scores)[0]
# Find the span with the maximum sum of scores for start and end indices.
total_scores = (tf.expand_dims(start_scores, axis=3) +
tf.expand_dims(end_scores, axis=2))
# Mask out scores where start_index > end_index.
start_idx = tf.reshape(tf.range(max_num_tokens), [1, 1, -1, 1])
end_idx = tf.reshape(tf.range(max_num_tokens), [1, 1, 1, -1])
invalid_index_mask = tf.tile((start_idx > end_idx),
[batch_size, max_num_slots, 1, 1])
total_scores = tf.where(invalid_index_mask,
tf.zeros_like(total_scores), total_scores)
max_span_index = tf.argmax(tf.reshape(
total_scores, [-1, max_num_slots, max_num_tokens**2]),
axis=-1)
span_start_index = tf.floordiv(max_span_index, max_num_tokens)
span_end_index = tf.floormod(max_span_index, max_num_tokens)
predictions["noncat_slot_start"] = span_start_index
predictions["noncat_slot_end"] = span_end_index
# Add inverse alignments.
predictions["noncat_alignment_start"] = features[
"noncat_alignment_start"]
predictions["noncat_alignment_end"] = features["noncat_alignment_end"]
return predictions
def rotate_pano_horizontally(input_feature_map, yaw_angle):
"""Rotates input_feature_map by yaw_angle by horizontally translating pixels.
The layer is differentiable with respect to yaw_angle and input_feature_map.
yaw_angle is positive for CCW rotation about the z-axis where the coordinates
are constructed with z-axis facing up.
Args:
input_feature_map: panoramic image or neural feature maps of shape [B, H, W,
C].
yaw_angle: A tensor of shape `[B]` which represents the desired rotation of
input_feature_map. yaw_angle is in units of radians. A positive yaw_angle
rolls pixels left.
Returns:
A rotated feature map with dimensions `[B, H, W, C]`
Reference:
[1]: 'Spatial Transformer Networks', Jaderberg et. al,
(https://arxiv.org/abs/1506.02025)
"""
# Number of input dimensions.
tfshape = tf.shape(input_feature_map)
batch_size = tfshape[0]
height = tfshape[1]
width = tfshape[2]
float32_width = tf.cast(width, dtype=tf.float32)
float32_height = tf.cast(height, dtype=tf.float32)
x_offset = (yaw_angle / 2 / np.pi) * float32_width
x_grid = tf.linspace(0., float32_width - 1, width) # (W)
# 0.5 * original_image_width to match the convention described in comment
x_pixel_coord = x_grid[tf.newaxis] + x_offset[:, tf.newaxis] # (B, W)
x_pixel_coord = tf.tile(x_pixel_coord[:, tf.newaxis, :],
[1, height, 1]) # (B, H, W)
y_pixel_coord = tf.linspace(0., float32_height - 1,
height)[tf.newaxis, :, tf.newaxis] # (1, H, 1)
y_pixel_coord = tf.tile(y_pixel_coord, [batch_size, 1, width])
wrapped_x_pixel_coord = tf.floormod(x_pixel_coord, float32_width)
# Because these are panoramas, we can concatenate the first column to the
# right side. This allows us to interpolate values for coordinates that
# correspond to pixels that connects the left and right edges of the
# panorama.
input_feature_map = tf.concat(
[input_feature_map, input_feature_map[:, :, :1]], axis=2)
return resampler.resampler(
input_feature_map,
tf.stack([wrapped_x_pixel_coord, y_pixel_coord], axis=-1))
def ae_latent_softmax(latents_pred, latents_discrete, hparams):
"""Latent prediction and loss."""
vocab_size = 2**hparams.z_size
if hparams.num_decode_blocks < 2:
latents_logits = tf.layers.dense(latents_pred,
vocab_size,
name="extra_logits")
if hparams.logit_normalization:
latents_logits *= tf.rsqrt(
1e-8 + tf.reduce_mean(tf.square(latents_logits)))
loss = None
if latents_discrete is not None:
if hparams.soft_em:
# latents_discrete is actually one-hot of multinomial samples
assert hparams.num_decode_blocks == 1
loss = tf.nn.softmax_cross_entropy_with_logits_v2(
labels=latents_discrete, logits=latents_logits)
else:
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=latents_discrete, logits=latents_logits)
sample = multinomial_sample(latents_logits, vocab_size,
hparams.sampling_temp)
return sample, loss
# Multi-block case.
vocab_bits = int(math.log(vocab_size, 2))
assert vocab_size == 2**vocab_bits
assert vocab_bits % hparams.num_decode_blocks == 0
block_vocab_size = 2**(vocab_bits // hparams.num_decode_blocks)
latents_logits = [
tf.layers.dense(latents_pred,
block_vocab_size,
name="extra_logits_%d" % i)
for i in range(hparams.num_decode_blocks)
]
loss = None
if latents_discrete is not None:
losses = []
for i in range(hparams.num_decode_blocks):
d = tf.floormod(tf.floordiv(latents_discrete, block_vocab_size**i),
block_vocab_size)
losses.append(
tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=d, logits=latents_logits[i]))
loss = sum(losses)
samples = [
multinomial_sample(l, block_vocab_size, hparams.sampling_temp)
for l in latents_logits
]
sample = sum([s * block_vocab_size**i for i, s in enumerate(samples)])
return sample, loss
def unravel_index_2d(indices, dims):
"""Unravel index, for 2D inputs only.
See Numpy's unravel.
Args:
indices: <int32> [num_elements], coordinates into 2D row-major tensor.
dims: (N, M), dimensions of the 2D tensor.
Returns:
coordinates: <int32> [2, num_elements], row (1st) and column (2nd) indices.
"""
row_inds = tf.floordiv(indices, dims[1])
col_inds = tf.floormod(indices, dims[1])
return tf.stack([row_inds, col_inds], axis=0)
def int_to_bit(self, x_int, num_bits, base=2):
"""Turn x_int representing numbers into a bitwise (lower-endian) tensor.
Args:
x_int: Tensor containing integer to be converted into base
notation.
num_bits: Number of bits in the representation.
base: Base of the representation.
Returns:
Corresponding number expressed in base.
"""
x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
# pylint: disable=g-complex-comprehension
x_labels = [
tf.floormod(tf.floordiv(tf.to_int32(x_l),
tf.to_int32(base)**i), tf.to_int32(base))
for i in range(num_bits)
]
res = tf.concat(x_labels, axis=-1)
return tf.to_float(res)
def should_log(params):
"""Returns a Boolean `tf.Tensor` dictating whether we should log values."""
global_step = tf.train.get_or_create_global_step()
first_run = tf.equal(global_step, 1)
log_every = tf.equal(tf.floormod(global_step, params.log_every), 0)
return tf.logical_or(first_run, log_every)
def mix_data(example):
"""Function to mix the different datasets according to a schedule."""
del example
# This block computes the probability of mixing the primary task with
# the secondary tasks. 0 = only the primary task, 1 = only the secondary
# tasks.
if hparams.multiproblem_mixing_schedule == MixingSchedule.EXPONENTIAL:
prob = get_exp_sched_prob()
prob = tf.cond(
tf.equal(
tf.floormod(problem_step,
tf.cast(5e6, dtype=tf.int64)), 0),
lambda: tf.Print(prob, [prob], message="Probability"),
lambda: prob)
elif hparams.multiproblem_mixing_schedule == MixingSchedule.CONSTANT:
prob = get_const_sched_prob()
elif hparams.multiproblem_mixing_schedule == MixingSchedule.PRETRAIN:
prob = get_pretrain_sched_prob()
else:
raise ValueError("Unknown schedule %s" %
str(hparams.multiproblem_mixing_schedule))
tf.logging.info("Using the %s schedule to "
"train the MultiProblem." %
str(hparams.multiproblem_mixing_schedule))
tf.logging.info("Schedule mixing threshold "
"%.2f" %
hparams.multiproblem_schedule_threshold)
# If per-task thresholds are specified, use them.
thresholds = None
if hparams.multiproblem_per_task_threshold:
thresholds = hparams.multiproblem_per_task_threshold.split(
",")
thresholds = [float(t)
for t in thresholds] # Convert to floats.
thresholds_sum = sum(thresholds)
tf.logging.info("Per-task thresholds: %s." %
str(thresholds))
thresholds = [t / thresholds_sum
for t in thresholds] # Normalize.
thresholds = [
sum(thresholds[:i + 1]) for i in range(len(thresholds))
]
tf.logging.info("Per-task threshold sums: %s." %
str(thresholds))
if len(thresholds) != len(self.task_list):
tf.logging.warn(
"Specified %d thresholds but encountered %d tasks."
% (len(thresholds), len(self.task_list)))
thresholds = None
def sample_task(curr_task, num_tasks_left, randnum):
"""A recursive function to sample a task.
This function treats the probability as the threshold for the primary
task and divides the remaining probability mass across the other
tasks.
Args:
curr_task: The index of the task being considered for sampling.
num_tasks_left: Number of tasks remaining to possibly sample from.
randnum: The random number used to select the dataset.
Returns:
A Tensor representing an example from the task that was sampled
from.
"""
if num_tasks_left == 0:
return get_next_from_dataset(
dataset_iterators[curr_task])
if thresholds is not None: # Use per-task thresholds if specified.
prob_sum = thresholds[curr_task]
return tf.cond(
randnum < prob_sum, lambda: get_next_from_dataset(
dataset_iterators[curr_task]),
lambda: sample_task(curr_task + 1, num_tasks_left -
1, randnum))
# When curr_task is 0, the primary task, the new prob is the same as
# the original probability. `tf.greater` indicates that the primary
# task receives (1-prob) of the probability mass.
# Otherwise, `prob` is divided equally amongst all the secondary
# tasks.
new_prob = prob - (curr_task * prob /
(len(self.task_list) - 1))
return tf.cond(
tf.greater(randnum, new_prob),
lambda: get_next_from_dataset(dataset_iterators[
curr_task]), lambda: sample_task(
curr_task + 1, num_tasks_left - 1, randnum))
return tf.data.Dataset.from_tensors(
sample_task(0,
len(self.task_list) - 1,
tf.random_uniform([])))
def project_inner(self, index):
"""Projects an index with the same index set onto the inner components."""
return IndexMap(indices=tf.floormod(index.indices,
self.inner_index.num_segments),
num_segments=self.inner_index.num_segments,
batch_dims=index.batch_dims)
