def initialize(self, name=None): """ Performs helper initialization (to get initial state) """ with ops.name_scope(name, 'CustomHelperInitialize'): finished = gen_math_ops.equal(0, self._sequence_length) all_finished = math_ops.reduce_all(finished) initial_candidates = self._candidate_tas.read(0) def training_inputs(): """ Returns the training initial input """ embed_op = self._order_embedding_fn(self._input_tas.read(0)) with ops.control_dependencies([embed_op]): return array_ops.identity(embed_op) def start_inputs(): """ Returns the GO_ID initial input """ embed_op = self._order_embedding_fn(self._start_inputs) with ops.control_dependencies([embed_op]): return array_ops.identity(embed_op) # Getting initial inputs initial_inputs = control_flow_ops.case( [(gen_math_ops.equal(self._decoder_type, TRAINING_DECODER), training_inputs), (gen_math_ops.equal(self._decoder_type, GREEDY_DECODER), start_inputs), (gen_math_ops.equal(self._decoder_type, SAMPLE_DECODER), start_inputs)], default=training_inputs) next_inputs = \ control_flow_ops.cond(all_finished, lambda: self._zero_inputs, lambda: CandidateInputs( inputs=self._input_layer(initial_inputs), candidates=initial_candidates, candidates_emb=self._candidate_embedding_fn(initial_candidates))) return (finished, next_inputs)
def seeded_random(seeds, offset, shape, dtype, seed=None, name=None): """ Outputs random values from a uniform distribution. The random values are deterministic given a seed. :param seeds: A vector of seeds (Size: [batch,]) - If 0, defaults to seed attr, then graph seed, then random. :param offset: Integer to add to the seed to get a deterministic mask. :param shape: The shape required for each seed (e.g. [3, 5] with a batch of 10 will return [10, 3, 5]). :param dtype: The type of the output. `float16`, `float32`, `float64` :param seed: A Python integer. Used to create a default seed for the operation. :param name: A name for the operation (optional). :return: A tensor of the specified shape filled with deterministic random values. """ if dtype not in (dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64): raise ValueError('Invalid dtype %r' % dtype) with ops.name_scope(name, 'seeded_random', [shape]): seeds = ops.convert_to_tensor(seeds, dtype=dtypes.int32, name='seeds') shape = ops.convert_to_tensor(shape, dtype=dtypes.int32, name='shape') offset = ops.convert_to_tensor(offset, dtype=dtypes.int32, name='offset') size = math_ops.reduce_prod(shape) graph_seed, op_seed = random_seed.get_seed(seed) matrix_output = SEEDED_RANDOM_SO.seeded_random(seeds, offset, size, seed=graph_seed, seed2=op_seed) output = gen_array_ops.reshape( matrix_output, array_ops.concat([(-1, ), shape], axis=0)) return math_ops.cast(output, dtype)
def next_inputs(self, time, outputs, state, sample_ids, name=None): """ Computes the next inputs at a time step """ with ops.name_scope(name, 'CustomHelperNextInputs', [time, outputs, state, sample_ids]): next_time = time + 1 finished = (next_time >= self._sequence_length) all_finished = math_ops.reduce_all(finished) def get_next_inputs(): """ Retrieves the inputs for the next time step """ def get_training_inputs(): """ Selecting training inputs """ read_op = self._input_tas.read(next_time) with ops.control_dependencies([read_op]): return array_ops.identity(read_op) def get_sample_inputs(): """ Selecting greedy/sample inputs """ return sample_ids inputs_next_step = control_flow_ops.case( [(gen_math_ops.equal(self._decoder_type, TRAINING_DECODER), get_training_inputs), (gen_math_ops.equal(self._decoder_type, GREEDY_DECODER), get_sample_inputs), (gen_math_ops.equal(self._decoder_type, SAMPLE_DECODER), get_sample_inputs)], default=get_training_inputs) inputs_emb_next_step = self._input_layer( self._embedding_fn(inputs_next_step)) # Applying mask # inputs_one_hot: (b, 1, VOC, 1) # mask_t: (b, 1, VOC, VOC) # next_mask: (b, VOC) -- DenseTensor inputs_one_hot = array_ops.one_hot(inputs_next_step, self.vocab_size)[:, None, :, None] mask_t = _slice_mask(self._mask, [-1, next_time, -1, -1], time_major=self._time_major) next_mask = sparse_ops.sparse_reduce_sum(inputs_one_hot * mask_t, axis=[1, 2]) next_mask = gen_math_ops.minimum(next_mask, 1.) next_mask.set_shape([None, self.vocab_size]) # Prevents this branch from executing eagerly with ops.control_dependencies( [inputs_emb_next_step, next_mask]): return MaskedInputs( inputs=array_ops.identity(inputs_emb_next_step), mask=array_ops.identity(next_mask)) next_inputs = control_flow_ops.cond( all_finished, true_fn=lambda: self._zero_inputs, false_fn=get_next_inputs) # Returning return (finished, next_inputs, state)
def _update_distributed_as_chief(self, version_step=None): """ Performs the gradient averaging, updates the variables, and the global step :param version_step: A variable that represents the model's version :return: The update operation to run Note: This method is called by the chief when synchronization is required. """ # Creating sync_token queue with ops.device(self._global_step.device), ops.name_scope(''): self._sync_token_queue = data_flow_ops.FIFOQueue( capacity=-1, dtypes=self._global_step.dtype.base_dtype, shapes=(), name='sync_token_q', shared_name='sync_token_q') # Applying grads, then adding tokens to queue with ops.control_dependencies([self._apply_grad_op]): tokens = gen_array_ops.fill([self._num_workers], self._global_step) sync_op = self._sync_token_queue.enqueue_many((tokens, )) # Waiting for token in queue (sync point) with ops.control_dependencies([sync_op]): token = self._sync_token_queue.dequeue() update_ops = [state_ops.assign(self._local_step, token)] # Increasing version step if version_step is not None: update_ops += [state_ops.assign_add(version_step, 1)] # Returning return control_flow_ops.group(*update_ops)
def sample(self, time, outputs, state, name=None): """ Samples the id for the next time step (or -1 for teacher forcing) """ with ops.name_scope(name, 'CustomHelperSample', [time, outputs, state]): def training(): """ Selecting training / teacher forcing """ fill_op = gen_array_ops.fill([array_ops.shape(outputs)[0]], -1) with ops.control_dependencies([fill_op]): return array_ops.identity(fill_op) def greedy(): """ Selecting greedy """ argmax_op = math_ops.argmax(outputs, axis=-1, output_type=dtypes.int32) with ops.control_dependencies([argmax_op]): return array_ops.identity(argmax_op) def sample(): """ Sampling """ logits = outputs if self._softmax_temperature is None else outputs / self._softmax_temperature sample_id_sampler = categorical.Categorical(logits=logits) sample_op = sample_id_sampler.sample(seed=self._seed) with ops.control_dependencies([sample_op]): return array_ops.identity(sample_op) return control_flow_ops.case( [(gen_math_ops.equal(self._decoder_type, TRAINING_DECODER), training), (gen_math_ops.equal(self._decoder_type, GREEDY_DECODER), greedy), (gen_math_ops.equal(self._decoder_type, SAMPLE_DECODER), sample)], default=training)
def zero_state(self, batch_size, dtype): """ Return an initial (zero) state tuple for this `IdentityCell`. :param batch_size: `0D` integer tensor: the batch size. :param dtype: The internal state data type. :return: A zeroed out scalar representing the initial state of the cell. """ with ops.name_scope(type(self).__name__ + 'ZeroState', values=[batch_size]): return array_ops.zeros([], dtype=dtypes.int32)
def zero_state(self, batch_size, dtype): """ Return an initial (zero) state tuple for this cell. :param batch_size: `0D` integer tensor: the batch size. :param dtype: The internal state data type. :return: A tuple containing zeroed out tensors and, possibly, empty TA objects. """ with ops.name_scope(type(self).__name__ + 'ZeroState', values=[batch_size]): return self._cell.zero_state(batch_size, dtype)
def zero_state(self, batch_size, dtype): """ Return an initial (zero) state tuple for this `ArrayConcatWrapper`. :param batch_size: `0D` integer tensor: the batch size. :param dtype: The internal state data type. :return: An `ArrayConcatWrapperState` tuple containing zeroed out tensors and, possibly, empty TA objects. """ with ops.name_scope(type(self).__name__ + 'ZeroState', values=[batch_size]): return ArrayConcatWrapperState( cell_state=self._cell.zero_state(batch_size, dtype), time=array_ops.zeros([], dtype=dtypes.int32))
def __init__(self, decoder_type, inputs, order_embedding, candidate_embedding, sequence_length, candidates, input_layer=None, time_major=False, softmax_temperature=None, seed=None, name=None): """ Constructor :param decoder_type: An uint8 representing TRAINING_DECODER, GREEDY_DECODER, or SAMPLE_DECODER :param inputs: The decoder input (b, dec_len) :param order_embedding: The order embedding vector :param candidate_embedding: The candidate embedding vector :param sequence_length: The length of each input (b,) :param candidates: The candidates at each time step -- Size: (b, nb_cand, max_candidates) :param input_layer: Optional. A layer to apply on the inputs :param time_major: If true indicates that the first dimension is time, otherwise it is batch size :param softmax_temperature: Optional. Softmax temperature. None, scalar, or size: (batch_size,) :param seed: Optional. The sampling seed :param name: Optional scope name. """ # pylint: disable=too-many-arguments with ops.name_scope(name, "CustomHelper", [inputs, sequence_length, order_embedding, candidate_embedding]): inputs = ops.convert_to_tensor(inputs, name="inputs") candidates = ops.convert_to_tensor(candidates, name="candidates") self._inputs = inputs self._order_embedding_fn = _get_embedding_fn(order_embedding) self._candidate_embedding_fn = _get_embedding_fn(candidate_embedding) if not time_major: inputs = nest.map_structure(_transpose_batch_time, inputs) candidates = nest.map_structure(_transpose_batch_time, candidates) self._input_tas = nest.map_structure(_unstack_ta, inputs) self._candidate_tas = nest.map_structure(_unstack_ta, candidates) self._decoder_type = decoder_type self._sequence_length = ops.convert_to_tensor(sequence_length, name="sequence_length") if self._sequence_length.get_shape().ndims != 1: raise ValueError("Expected vector for sequence_length. Shape: %s" % self._sequence_length.get_shape()) self._input_layer = input_layer if input_layer is not None else lambda x: x self._batch_size = array_ops.size(sequence_length) self._start_inputs = gen_array_ops.fill([self._batch_size], GO_ID) self._softmax_temperature = softmax_temperature self._seed = seed # Compute input shape self._zero_inputs = \ CandidateInputs(inputs= array_ops.zeros_like(self._input_layer(self._order_embedding_fn(self._start_inputs))), candidates=array_ops.zeros_like(candidates[0, :]), candidates_emb=array_ops.zeros_like(self._candidate_embedding_fn(candidates[0, :]))) # Preventing div by zero # Adding an extra dim to the matrix, so we can broadcast with the outputs shape if softmax_temperature is not None: self._softmax_temperature = gen_math_ops.maximum(1e-10, self._softmax_temperature) if self._softmax_temperature.get_shape().ndims == 1: self._softmax_temperature = self._softmax_temperature[:, None]
def step(self, time, inputs, state, name=None): """ Performs a decoding step :param time: scalar `int32` tensor. :param inputs: A (structure of) input tensors. (** This is a MaskedInputs tuple **) :param state: A (structure of) state tensors and TensorArrays. :param name: Name scope for any created operations. :return: (outputs, next_state, next_inputs, finished) """ assert isinstance( inputs, CandidateInputs), 'The inputs must be of type "CandidateInputs"' with ops.name_scope(name, "BasicDecoderStep", (time, inputs, state)): inputs, candidates, candidates_emb = inputs.inputs, inputs.candidates, inputs.candidates_emb cell_outputs, cell_state = self._cell(inputs, state) cell_state_output = cell_outputs # Corresponds to cell_state.h (before output layer) if self._output_layer is not None: cell_outputs = self._output_layer(cell_outputs) # Adding a bias dimension, then computing candidate logits and masking PAD_IDs cell_outputs = array_ops.pad(cell_outputs, [(0, 0), (0, 1)], constant_values=1.) cell_outputs = math_ops.reduce_sum(cell_outputs[:, None, :] * candidates_emb, axis=-1) output_mask = math_ops.cast(gen_math_ops.greater(candidates, 0), dtypes.float32) cell_outputs = gen_math_ops.add(cell_outputs, (1. - output_mask) * LARGE_NEGATIVE) # Sampling and computing next inputs sample_ids = self._helper.sample(time=time, outputs=(cell_outputs, candidates), state=cell_state) (finished, next_inputs, next_state) = self._helper.next_inputs(time=time, outputs=cell_outputs, state=cell_state, sample_ids=sample_ids) if self.extract_state: outputs = BasicDecoderWithStateOutput(cell_outputs, cell_state_output, sample_ids) else: outputs = seq2seq.BasicDecoderOutput(cell_outputs, sample_ids) return outputs, next_state, next_inputs, finished
def zero_state(self, batch_size, dtype): """ Return an initial (zero) state tuple for this `AttentionWrapper`. :param batch_size: `0D` integer tensor: the batch size. :param dtype: The internal state data type. :return: An `SelfAttentionWrapperState` tuple containing zeroed out tensors. """ with ops.name_scope(type(self).__name__ + 'ZeroState', values=[batch_size]): # Using batch_size * 0, rather than just 0 to have a dynamic dimension initial_cell_state = self._cell.zero_state(batch_size, dtype) initial_memory = array_ops.zeros( [batch_size, batch_size * 0, self._memory_size], dtype=self._dtype) return SelfAttentionWrapperState(cell_state=initial_cell_state, time=array_ops.zeros( [], dtype=dtypes.int32), memory=initial_memory)
def zero_state(self, batch_size, dtype): """ Return an initial (zero) state tuple for this `AttentionWrapper`. :param batch_size: `0D` integer tensor: the batch size. :param dtype: The internal state data type. :return: AttentionWrapperState` tuple containing zeroed out tensors and, possibly, empty `TensorArrays`. """ with ops.name_scope(type(self).__name__ + 'ZeroState', values=[batch_size]): if self._initial_cell_state is not None: cell_state = self._initial_cell_state else: cell_state = self._cell.zero_state(batch_size, dtype) error_message = ( 'When calling zero_state of AttentionWrapper %s: ' % self._base_name + 'Non-matching batch sizes between the memory encoder output) and the requested batch ' 'size. Are you using the BeamSearchDecoder? If so, make sure your encoder output has been ' 'tiled to beam_width via tf.contrib.seq2seq.tile_batch, and the batch_size= argument ' 'passed to zero_state is batch_size * beam_width.') with ops.control_dependencies( self._batch_size_checks(batch_size, error_message)): cell_state = nest.map_structure( lambda state: array_ops.identity( state, name='checked_cell_state'), cell_state) initial_alignments = [ attention_mechanism.initial_alignments(batch_size, dtype) for attention_mechanism in self._attention_mechanisms ] return AttentionWrapperState( cell_state=cell_state, time=array_ops.zeros([], dtype=dtypes.int32), attention=_zero_state_tensors(self._attention_layer_size, batch_size, dtype), alignments=self._item_or_tuple(initial_alignments), attention_state=self._item_or_tuple( attention_mechanism.initial_state(batch_size, dtype) for attention_mechanism in self._attention_mechanisms), alignment_history=self._item_or_tuple( tensor_array_ops.TensorArray(dtype, size=0, dynamic_size=True, element_shape=alignment.shape) if self._alignment_history else () for alignment in initial_alignments))
def zero_state(self, batch_size, dtype): """ Return an initial (zero) state tuple for this `IdentityCell`. :param batch_size: `0D` integer tensor: the batch size. :param dtype: The internal state data type. :return: A zeroed out scalar representing the initial state of the cell. """ with ops.name_scope(type(self).__name__ + 'ZeroState', values=[batch_size]): if self._feeder_cell is None: feeder_init_state = array_ops.zeros([], dtype=dtype) elif self._feeder_init_state is not None: feeder_init_state = self._feeder_init_state else: feeder_init_state = self._feeder_cell.zero_state( batch_size, dtype) # Empty past attentions if self._past_attns is None: head_size = self._emb_size // self._nb_heads past_attns_shape = [ batch_size, self._nb_layers, 2, self._nb_heads, 0 * batch_size, head_size ] self._past_attns = array_ops.zeros(past_attns_shape, dtype=dtypes.float32) # No Context - Returning a zero past attention if self._context is None: return TransformerCellState(past_attentions=self._past_attns, feeder_state=feeder_init_state, time=array_ops.zeros( [], dtype=dtypes.int32)) # Context provided - Computing attention by running a single block step _, present_attns, _ = self._step( inputs=self._context_word_embedding_fn(self._context), past_attns=self._past_attns, time=0, feeder_cell=None, feeder_state=None) return TransformerCellState(past_attentions=present_attns, feeder_state=feeder_init_state, time=array_ops.zeros( [], dtype=dtypes.int32))
def _update_standalone(self, version_step=None): """ Performs the gradient averaging, updates the variables, and the global step :param version_step: A variable that represents the model's version :return: The update operation to run Note: This method is called when there are no workers (no synchronization) """ with ops.device(self._global_step.device), ops.name_scope(''): with ops.control_dependencies([self._apply_grad_op]): update_ops = [ state_ops.assign(self._local_step, self._global_step) ] # Increasing version step if version_step is not None: update_ops += [state_ops.assign_add(version_step, 1)] # Returning return control_flow_ops.group(*update_ops)
def _update_distributed_as_worker(self): """ Performs the gradient averaging, updates the variables, and the global step :param version_step: A variable that represents the model's version :return: The update operation to run Note: This method is called by a worker when synchronization is required. """ # Creating sync_token queue with ops.device(self._global_step.device), ops.name_scope(''): self._sync_token_queue = data_flow_ops.FIFOQueue( capacity=-1, dtypes=self._global_step.dtype.base_dtype, shapes=(), name='sync_token_q', shared_name='sync_token_q') # Waiting for token in queue (sync point) token = self._sync_token_queue.dequeue() return state_ops.assign(self._local_step, token)
def next_inputs(self, time, outputs, state, sample_ids, name=None): """ Computes the next inputs at a time step """ with ops.name_scope(name, 'CustomHelperNextInputs', [time, outputs, state, sample_ids]): next_time = time + 1 finished = (next_time >= self._sequence_length) all_finished = math_ops.reduce_all(finished) def get_next_inputs(): """ Retrieves the inputs for the next time step """ def get_training_inputs(): """ Selecting training inputs """ read_op = self._input_tas.read(next_time) with ops.control_dependencies([read_op]): return array_ops.identity(read_op) def get_sample_inputs(): """ Selecting greedy/sample inputs """ return sample_ids inputs_next_step = control_flow_ops.case( [(gen_math_ops.equal(self._decoder_type, TRAINING_DECODER), get_training_inputs), (gen_math_ops.equal(self._decoder_type, GREEDY_DECODER), get_sample_inputs), (gen_math_ops.equal(self._decoder_type, SAMPLE_DECODER), get_sample_inputs)], default=get_training_inputs) inputs_emb_next_step = self._input_layer(self._order_embedding_fn(inputs_next_step)) candidate_next_step = self._candidate_tas.read(next_time) candidate_emb_next_step = self._candidate_embedding_fn(candidate_next_step) # Prevents this branch from executing eagerly with ops.control_dependencies([inputs_emb_next_step, candidate_next_step, candidate_emb_next_step]): return CandidateInputs(inputs=array_ops.identity(inputs_emb_next_step), candidates=array_ops.identity(candidate_next_step), candidates_emb=array_ops.identity(candidate_emb_next_step)) next_inputs = control_flow_ops.cond(all_finished, true_fn=lambda: self._zero_inputs, false_fn=get_next_inputs) # Returning return (finished, next_inputs, state)
def sample(self, time, outputs, state, name=None): """ Samples the id for the next time step (or -1 for teacher forcing) Note: outputs is a tuple of (cell_outputs, candidate) """ cell_outputs, candidate = outputs with ops.name_scope(name, 'CustomHelperSample', [time, outputs, state]): def training(): """ Selecting training / teacher forcing """ fill_op = gen_array_ops.fill([array_ops.shape(cell_outputs)[0]], -1) with ops.control_dependencies([fill_op]): return array_ops.identity(fill_op) def greedy(): """ Selecting greedy """ argmax_id = math_ops.cast(math_ops.argmax(cell_outputs, axis=-1), dtypes.int32) nb_candidate = array_ops.shape(candidate)[1] candidate_ids = \ math_ops.reduce_sum(array_ops.one_hot(argmax_id, nb_candidate, dtype=dtypes.int32) * candidate, axis=-1) with ops.control_dependencies([candidate_ids]): return array_ops.identity(candidate_ids) def sample(): """ Sampling """ logits = cell_outputs if self._softmax_temperature is None else cell_outputs / self._softmax_temperature sample_id_sampler = categorical.Categorical(logits=logits) sample_ids = sample_id_sampler.sample(seed=self._seed) nb_candidate = array_ops.shape(candidate)[1] reduce_op = math_ops.reduce_sum(array_ops.one_hot(sample_ids, nb_candidate, dtype=dtypes.int32) * candidate, axis=-1) with ops.control_dependencies([reduce_op]): return array_ops.identity(reduce_op) return control_flow_ops.case([(gen_math_ops.equal(self._decoder_type, TRAINING_DECODER), training), (gen_math_ops.equal(self._decoder_type, GREEDY_DECODER), greedy), (gen_math_ops.equal(self._decoder_type, SAMPLE_DECODER), sample)], default=training)
def step(self, time, inputs, state, name=None): """ Performs a decoding step :param time: scalar `int32` tensor. :param inputs: A (structure of) input tensors. (** This is a MaskedInputs tuple **) :param state: A (structure of) state tensors and TensorArrays. :param name: Name scope for any created operations. :return: (outputs, next_state, next_inputs, finished) """ assert isinstance( inputs, (MaskedInputs, ops.Tensor)), 'Expected "MaskedInputs" or a Tensor.' with ops.name_scope(name, "BasicDecoderStep", (time, inputs, state)): inputs, output_mask = inputs, None if isinstance(inputs, MaskedInputs): inputs, output_mask = inputs.inputs, inputs.mask cell_outputs, cell_state = self._cell(inputs, state) cell_state_output = cell_outputs # Corresponds to cell_state.h (before output layer) if self._output_layer is not None: cell_outputs = self._output_layer(cell_outputs) if output_mask is not None: cell_outputs = gen_math_ops.add( cell_outputs, (1. - output_mask) * LARGE_NEGATIVE) sample_ids = self._helper.sample(time=time, outputs=cell_outputs, state=cell_state) (finished, next_inputs, next_state) = self._helper.next_inputs(time=time, outputs=cell_outputs, state=cell_state, sample_ids=sample_ids) if self.extract_state: outputs = BasicDecoderWithStateOutput(cell_outputs, cell_state_output, sample_ids) else: outputs = seq2seq.BasicDecoderOutput(cell_outputs, sample_ids) return outputs, next_state, next_inputs, finished
def apply_gradients(self, grads_and_vars, global_step=None, name=None): """ Accumulates gradients for the variables and stores them in the accumulator Note: This does not update the variables, it just accumulate the gradients :param grads_and_vars: List of (gradient, variable) pairs as returned by compute_gradients() :param global_step: Optional variable to increment by one after the variables have been updated. :param name: Optional name for the returned op. Default to name passed to the optimizer constructor. :return: The training operation to be run by each replica Raises - ValueError if the grads_and_vars is empty - ValueError if global step is not provided - ValueError if update() has already been called """ # Making sure grads_and_var and global_step are provided if not grads_and_vars: raise ValueError('Must supply at least one variable.') if not global_step: raise ValueError('You must provide a global_step variable') if self._finalized: raise ValueError( 'The optimizer has been finalized. You cannot call this method after update().' ) train_ops = [] accumulated_grad = [] var_list = [] chief_init_ops = [] self._global_step = global_step # Colocating local step to prevent it from being placed on the parameter server local_anchor = gen_control_flow_ops.no_op() with ops.device(local_anchor.device): self._local_step = variable_scope.variable( initial_value=0, trainable=False, collections=[ops.GraphKeys.LOCAL_VARIABLES], dtype=global_step.dtype.base_dtype, name='sync_rep_local_step') # Setting initial step self.local_step_init_op = state_ops.assign(self._local_step, global_step) chief_init_ops += [self.local_step_init_op] with ops.name_scope(None, self._name): # Creating accumulators for grad, var in grads_and_vars: var_list += [var] with ops.device(var.device): # No gradient - Pass-Through if grad is None: accumulated_grad += [None] continue # Sparse Variable - Accumulating over the dense shape elif isinstance(grad, ops.IndexedSlices): grad_accum = self._create_sparse_accumulator(var, grad) train_ops += [ grad_accum.apply_indexed_slices_grad( grad, local_step=self._local_step) ] accumulated_grad += [ self._take_sparse_grad(grad_accum, grad) ] chief_init_ops += [ grad_accum.set_global_step(global_step, name='SetGlobalStep') ] # Dense Variable elif isinstance(grad, ops.Tensor): grad_accum = self._create_dense_accumulator(var, grad) train_ops += [ grad_accum.apply_grad(grad, local_step=self._local_step) ] accumulated_grad += [ self._take_dense_grad(grad_accum, grad) ] chief_init_ops += [ grad_accum.set_global_step(global_step, name='SetGlobalStep') ] # Unknown else: raise RuntimeError('Unsupported gradient type.') # Building update_op with ops.device(self._global_step.device), ops.name_scope(''): accumulated_grad = [ ensure_finite(gradient) for gradient in accumulated_grad ] if self._max_gradient_norm: accumulated_grad, _ = clip_ops.clip_by_global_norm( accumulated_grad, self._max_gradient_norm) self._apply_grad_op = self._optimizer.apply_gradients( zip(accumulated_grad, var_list), global_step) # Building chief init ops self.chief_init_op = control_flow_ops.group(*chief_init_ops) # Building train_op return control_flow_ops.group(*train_ops)
def seeded_dropout(inputs, seeds, keep_probs, offset=None, noise_shape=None, seed=None, name=None): """ Computes dropout (with a deterministic mask). Every item in the batch has a deterministic seed to compute the deterministic mask With probability `keep_probs`, outputs the input element scaled up by `1 / keep_prob`, otherwise outputs `0`. The scaling is so that the expected sum is unchanged. By default, each element is kept or dropped independently. If `noise_shape` is specified, it must be broadcastable to the shape of `x`, and only dimensions with `noise_shape[i] == shape(x)[i]` will make independent decisions. For example, if `shape(x) = [k, l, m, n]` and `noise_shape = [k, 1, 1, n]`, each batch and channel component will be kept independently and each row and column will be kept or not kept together. :param inputs: A floating point tensor. :param seeds: A tensor representing the seed for each item in the batch. (Size: (batch,)) :param keep_probs: A scalar or vector of size (batch,). The probability that each element is kept. :param offset: Integer. Alternative offset to apply to compute the deterministic mask (e.g. in a loop). :param noise_shape: A 1-D `Tensor` of type `int32`, represents the shape for randomly generated keep/drop flags. :param seed: A Python integer. Used to create a default seed for the operation. :param name: name: A name for this operation (optional). :return: A Tensor of the same shape of `x`. """ if offset is None: seeded_dropout.offset += 40555607 # If inputs is a scalar, this is likely the 'time' attribute in a state, we don't want to mask it # Same thing for integers - We can safely ignore them # So we don't want to mask it if not inputs.shape or inputs.dtype.is_integer: return inputs with ops.name_scope(name, 'seeded_dropout', [inputs]): inputs = ops.convert_to_tensor(inputs, name='x') if not inputs.dtype.is_floating: raise ValueError( 'Expected a floating point tensor. Got a %s tensor instead.' % inputs.dtype) if isinstance(keep_probs, float) and not 0 < keep_probs <= 1: raise ValueError( 'keep_probs must be a scalar tensor or a float in the range (0, 1], got %g' % keep_probs) # Early return if nothing needs to be dropped. if isinstance(keep_probs, float) and keep_probs == 1: return inputs # Not supported in eager mode if context.executing_eagerly(): raise ValueError('This function is not supported in eager mode.') # Converting to tensor keep_probs = ops.convert_to_tensor(keep_probs, dtype=inputs.dtype, name='keep_probs') keep_probs = gen_math_ops.maximum(0., gen_math_ops.minimum(1., keep_probs)) keep_probs = gen_array_ops.reshape(keep_probs, [-1] + [1] * (len(inputs.shape) - 1)) all_keep_probs_are_one = math_ops.reduce_all( gen_math_ops.equal(keep_probs, 1.)) # Computing noise shape noise_shape = nn_ops._get_noise_shape(inputs, noise_shape) # pylint: disable=protected-access def get_dropout_mask(): """ Computes the dropout mask """ # random_tensor = uniform [keep_probs, 1.0 + keep_probs) random_tensor = keep_probs random_tensor += seeded_random( seeds, offset=offset if offset is not None else seeded_dropout.offset, shape=noise_shape[1:], dtype=inputs.dtype, seed=seed) # 0. if [keep_probs, 1.0) and 1. if [1.0, 1.0 + keep_prob) binary_tensor = gen_math_ops.floor(random_tensor) ret = math_ops.divide(inputs, keep_probs) * binary_tensor ret.set_shape(inputs.get_shape()) # Setting control flow ops to avoid computing this function if not required with ops.control_dependencies([ret]): return array_ops.identity(ret) # Returning the dropout mask return control_flow_ops.cond(all_keep_probs_are_one, true_fn=lambda: inputs, false_fn=get_dropout_mask)
def __init__(self, decoder_type, inputs, embedding, sequence_length, mask, input_layer=None, time_major=False, softmax_temperature=None, seed=None, name=None): """ Constructor :param decoder_type: An uint8 representing TRAINING_DECODER, GREEDY_DECODER, or SAMPLE_DECODER :param inputs: The decoder input (b, dec_len) :param embedding: The embedding vector :param sequence_length: The length of each input (b,) :param mask: [SparseTensor] Mask to apply at each time step -- Size: (b, dec_len, vocab_size, vocab_size) :param input_layer: Optional. A layer to apply on the inputs :param time_major: If true indicates that the first dimension is time, otherwise it is batch size :param softmax_temperature: Optional. Softmax temperature. None or size: (batch_size,) :param seed: Optional. The sampling seed :param name: Optional scope name. """ # pylint: disable=too-many-arguments with ops.name_scope(name, "CustomHelper", [inputs, sequence_length, embedding]): assert isinstance(mask, SparseTensor), 'The mask must be a SparseTensor' inputs = ops.convert_to_tensor(inputs, name="inputs") self._inputs = inputs self._mask = mask self._time_major = time_major self._embedding_fn = embedding if callable( embedding) else lambda ids: embedding_lookup(embedding, ids) if not time_major: inputs = nest.map_structure(_transpose_batch_time, inputs) self._input_tas = nest.map_structure(_unstack_ta, inputs) self._decoder_type = decoder_type self._sequence_length = ops.convert_to_tensor( sequence_length, name="sequence_length") if self._sequence_length.get_shape().ndims != 1: raise ValueError( "Expected vector for sequence_length. Shape: %s" % self._sequence_length.get_shape()) self._input_layer = input_layer if callable( input_layer) else lambda x: x self._batch_size = array_ops.size(sequence_length) self._start_inputs = gen_array_ops.fill([self._batch_size], GO_ID) self._softmax_temperature = softmax_temperature self._seed = seed self.vocab_size = VOCABULARY_SIZE self._zero_inputs = \ MaskedInputs(inputs=array_ops.zeros_like(self._input_layer(self._embedding_fn(self._start_inputs))), mask=_slice_mask(self._mask, slicing=[-1, 0, GO_ID, -1], squeeze=True, time_major=self._time_major)) # Preventing div by zero # Adding an extra dim to the matrix, so we can broadcast with the outputs shape if softmax_temperature is not None: self._softmax_temperature = gen_math_ops.maximum( 1e-10, self._softmax_temperature) if self._softmax_temperature.get_shape().ndims == 1: self._softmax_temperature = self._softmax_temperature[:, None]