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)