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 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))
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 call(self, inputs, state): # pylint: disable=arguments-differ """ Perform a step of attention-wrapped RNN :param inputs: (Possibly nested tuple of) Tensor, the input at this time step. :param state: An instance of `SelfAttentionWrapperState` containing tensors from the previous time step. :return: A tuple `(attention_or_cell_output, next_state)`, where: - `attention_or_cell_output` depending on `output_attention`. - `next_state` is an instance of `SelfAttentionWrapperState` containing the state calculated at this time step. """ if not isinstance(state, SelfAttentionWrapperState): raise TypeError( 'Expected state to be instance of AttentionWrapperState. Received type %s instead.' % type(state)) # Getting batch size batch_size = array_ops.shape(inputs)[0] assert len(inputs.shape) == 2, 'Expected inputs to be of rank 2' def get_next_memory_and_attn(): """ Gets the next memory and attention """ next_memory = array_ops.concat( [ state.memory, # [b, t, mem_size] array_ops.expand_dims(self._input_fn(inputs), axis=1) ], axis=1) next_attention = self._compute_attention(inputs, next_memory) with ops.control_dependencies([next_memory, next_attention]): return array_ops.identity(next_memory), array_ops.identity( next_attention) def get_zero_memory_and_attn(): """ Time = 0, we don't concatenate to memory and attention is all 0. """ next_memory = state.memory next_attention = array_ops.zeros( [batch_size, self._attention_layer_size], dtype=inputs.dtype) with ops.control_dependencies([next_memory, next_attention]): return array_ops.identity(next_memory), array_ops.identity( next_attention) # Computing memory and attention memory, attention = control_flow_ops.cond( gen_math_ops.equal(state.time, 0), true_fn=get_zero_memory_and_attn, false_fn=get_next_memory_and_attn) # Calculate the true inputs to the cell based on the previous attention value. cell_inputs = self._cell_input_fn(inputs, attention) cell_state = state.cell_state cell_output, cell_state = self._cell(cell_inputs, cell_state) # Extracting computed context next_state = SelfAttentionWrapperState(cell_state=cell_state, time=state.time + 1, memory=memory) # Returning cell output or attention if self._output_attention: return attention, next_state return cell_output, next_state
def _take_sparse_grad(grad_accum, grad): """ Computes the gradient for a SparseConditionalAccumulator :param grad_accum: The gradient accumulator where gradients are stored :param grad: An instance of the gradient stored in the accumulator :return: The avg gradient to apply (or a zero-like object if no gradients are stored) :type grad_accum: data_flow_ops.SparseConditionalAccumulator """ def _take_grad(): """ Computes the gradient from the accumulator """ avg_grad = grad_accum.take_indexed_slices_grad(num_required=1) with ops.control_dependencies([avg_grad]): return ops.IndexedSlices(values=array_ops.identity( avg_grad.values), indices=avg_grad.indices, dense_shape=avg_grad.dense_shape) def _zero_grad(): """ Returns a zeroed-out gradient """ zero_values = array_ops.zeros_like(grad.values) with ops.control_dependencies([zero_values]): return ops.IndexedSlices( values=array_ops.identity(zero_values), indices=math_ops.cast(grad.indices, dtypes.int64), dense_shape=math_ops.cast(grad.dense_shape, dtypes.int64)) return control_flow_ops.cond(gen_math_ops.equal( grad_accum.num_accumulated(), 0), true_fn=_zero_grad, false_fn=_take_grad)
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))
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 initialize(self): """ Initialize the beam helper - Called in beam_decoder.initialize() :return: `(finished, start_inputs, initial_cell_state)`. """ finished, zero_inputs, zero_mask = self._finished, self._zero_inputs, self._zero_mask all_finished = math_ops.reduce_all( gen_math_ops.equal(0, self._sequence_length)) initial_inputs = self._embedding_fn(self._start_tokens) # Start Inputs start_inputs = control_flow_ops.cond( all_finished, lambda: zero_inputs, lambda: MaskedInputs(inputs=self._split_batch_beams( self._input_layer(initial_inputs), self._input_size), mask=zero_mask)) # Returning return finished, start_inputs, self._initial_cell_state
def initialize(self): """ Initialize the beam helper - Called in beam_decoder.initialize() :return: `(finished, start_inputs, initial_cell_state)`. """ finished, zero_inputs = self._finished, self._zero_inputs all_finished = math_ops.reduce_all(gen_math_ops.equal(0, self._sequence_length)) initial_inputs = self._order_embedding_fn(self._start_tokens) initial_candidates = self._candidate_tas.read(0) # Start Inputs start_inputs = control_flow_ops.cond(all_finished, lambda: zero_inputs, lambda: CandidateInputs( inputs=self._split_batch_beams(self._input_layer(initial_inputs), self._input_size), candidates=initial_candidates, candidates_emb=self._candidate_embedding_fn(initial_candidates))) return finished, start_inputs, self._initial_cell_state
def _take_dense_grad(grad_accum, grad): """ Computes the gradient for a ConditionalAccumulator :param grad_accum: The gradient accumulator where gradients are stored :param grad: An instance of the gradient stored in the accumulator :return: The avg gradient to apply (or a zero-like object if no gradients are stored) :type grad_accum: data_flow_ops.ConditionalAccumulator """ def _take_grad(): """ Computes the gradient from the accumulator """ avg_grad = grad_accum.take_grad(num_required=1) with ops.control_dependencies([avg_grad]): return array_ops.identity(avg_grad) def _zero_grad(): """ Returns a zeroed-out gradient """ zero_like_grad = array_ops.zeros_like(grad) with ops.control_dependencies([zero_like_grad]): return array_ops.identity(zero_like_grad) return control_flow_ops.cond(gen_math_ops.equal( grad_accum.num_accumulated(), 0), true_fn=_zero_grad, false_fn=_take_grad)
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)