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)