def _add_ml_loss(self, is_training): logits = self._add_decoder(is_training=is_training, decoder_inputs=self.decoder_inputs) with self.graph.as_default(): self.preds = tf.to_int32(tf.argmax( logits, axis=-1)) # shape: (batch_size, max_timestep) self.istarget = tf.to_float(tf.not_equal( self.xy, 0)) # shape: (batch_size, max_timestep) self.acc_full = tf.reduce_sum( tf.to_float(tf.equal(self.preds, self.xy)) * self.istarget) / (tf.reduce_sum(self.istarget)) self.acc_sum = tf.reduce_sum( tf.to_float( tf.equal(self.preds[hp.article_maxlen + 1:], self.xy[hp.article_maxlen + 1:])) * self.istarget) / (tf.reduce_sum(self.istarget)) self.rouge = tf.reduce_sum( rouge_l_fscore(self.preds[hp.article_maxlen + 1:], self.xy[hp.article_maxlen + 1:])) / float( hp.batch_size) ml_loss = -100 if is_training: # Loss self.y_smoothed = label_smoothing( tf.one_hot(self.xy, depth=self.vocab_size)) loss = tf.nn.softmax_cross_entropy_with_logits_v2( logits=self.logits, labels=self.y_smoothed) ml_loss = tf.reduce_sum(loss * self.istarget, name='fake_ml_loss') / (tf.reduce_sum( self.istarget)) return ml_loss
def _add_rl_loss(self): sample_logits, sample_preds = self._rl_autoinfer(greedy=tf.constant( value=False, dtype=tf.bool), name='sample_loop') greedy_logits, greedy_preds = self._rl_autoinfer(greedy=tf.constant( value=True, dtype=tf.bool), name='greedy_loop') self.sl = sample_logits sample_logits = tf.Print(input_=sample_logits, data=[sample_logits], message='sample_logits: ') greedy_logits = tf.Print(input_=greedy_logits, data=[greedy_logits], message='greedy_logits: ') self.reward_diff = tf.zeros(shape=()) for sent_i, ref in enumerate(tf.unstack(self.y)): real_y = ref[:tf.reduce_sum(tf.to_int32(tf.not_equal( self.y, 0)))] # remove the <PAD> in the end self.reward_diff += rouge_l_fscore( [greedy_preds[sent_i]], [real_y]) - rouge_l_fscore( [sample_preds[sent_i]], [real_y]) self.reward_diff = tf.Print(input_=self.reward_diff, data=[self.reward_diff], message='reward_diff: ') self.clipped_reward_diff = tf.math.minimum(x=self.reward_diff, y=tf.ones(shape=()) * hp.max_reward_diff) self.clipped_reward_diff = tf.math.maximum(x=self.clipped_reward_diff, y=-tf.ones(shape=()) * hp.max_reward_diff) # rl_loss = tf.reduce_sum(self.reward_diff * sample_logits) / (hp.batch_size * hp.summary_maxlen) rl_loss = tf.reduce_sum(self.clipped_reward_diff * sample_logits) / ( hp.batch_size * hp.summary_maxlen) rl_loss = tf.Print(input_=rl_loss, data=[rl_loss], message='rl_loss: ') return rl_loss # masked
def attention_decoder(_hps, v_size, _max_art_oovs, _enc_batch_extend_vocab, emb_dec_inputs, target_batch, _dec_in_state, _enc_states, enc_padding_mask, dec_padding_mask, cell, embedding, sampling_probability, alpha, unk_id, initial_state_attention=False, pointer_gen=True, use_coverage=False, prev_coverage=None, prev_decoder_outputs=[], prev_encoder_es = []): """ Args: _hps: parameter of the models. v_size: vocab size. _max_art_oovs: size of the oov tokens in current batch. _enc_batch_extend_vocab: encoder extended vocab batch. emb_dec_inputs: A list of 2D Tensors [batch_size x emb_dim]. target_batch: The indices of the target words. shape (max_dec_steps, batch_size) _dec_in_state: 2D Tensor [batch_size x cell.state_size]. _enc_states: 3D Tensor [batch_size x max_enc_steps x attn_size]. enc_padding_mask: 2D Tensor [batch_size x max_enc_steps] containing 1s and 0s; indicates which of the encoder locations are padding (0) or a real token (1). dec_padding_mask: 2D Tensor [batch_size x max_dec_steps] containing 1s and 0s; indicates which of the decoder locations are padding (0) or a real token (1). cell: rnn_cell.RNNCell defining the cell function and size. embedding: embedding matrix [vocab_size, emb_dim]. sampling_probability: sampling probability for scheduled sampling. alpha: soft-argmax argument. initial_state_attention: Note that this attention decoder passes each decoder input through a linear layer with the previous step's context vector to get a modified version of the input. If initial_state_attention is False, on the first decoder step the "previous context vector" is just a zero vector. If initial_state_attention is True, we use _dec_in_state to (re)calculate the previous step's context vector. We set this to False for train/eval mode (because we call attention_decoder once for all decoder steps) and True for decode mode (because we call attention_decoder once for each decoder step). pointer_gen: boolean. If True, calculate the generation probability p_gen for each decoder step. use_coverage: boolean. If True, use coverage mechanism. prev_coverage: If not None, a tensor with shape (batch_size, max_enc_steps). The previous step's coverage vector. This is only not None in decode mode when using coverage. prev_decoder_outputs: if not empty, a tensor of (len(prev_decoder_steps), batch_size, hidden_dim). The previous decoder output used for calculating the intradecoder attention during decode mode prev_encoder_es: if not empty, a tensor of (len(prev_encoder_es), batch_size, hidden_dim). The previous attention vector used for calculating the temporal attention during decode mode. Returns: outputs: A list of the same length as emb_dec_inputs of 2D Tensors of shape [batch_size x cell.output_size]. The output vectors. state: The final state of the decoder. A tensor shape [batch_size x cell.state_size]. attn_dists: A list containing tensors of shape (batch_size,max_enc_steps). The attention distributions for each decoder step. p_gens: List of length emb_dim, containing tensors of shape [batch_size, 1]. The values of p_gen for each decoder step. Empty list if pointer_gen=False. coverage: Coverage vector on the last step computed. None if use_coverage=False. vocab_scores: vocab distribution. final_dists: final output distribution. samples: contains sampled tokens. greedy_search_samples: contains greedy tokens. temporal_e: contains temporal attention. """ with variable_scope.variable_scope("attention_decoder") as scope: batch_size = _enc_states.get_shape()[0] # if this line fails, it's because the batch size isn't defined attn_size = _enc_states.get_shape()[2] # if this line fails, it's because the attention length isn't defined emb_size = emb_dec_inputs[0].get_shape()[1] # if this line fails, it's because the embedding isn't defined decoder_attn_size = _dec_in_state.c.get_shape()[1] tf.logging.info("batch_size %i, attn_size: %i, emb_size: %i", batch_size, attn_size, emb_size) # Reshape _enc_states (need to insert a dim) _enc_states = tf.expand_dims(_enc_states, axis=2) # now is shape (batch_size, max_enc_steps, 1, attn_size) # To calculate attention, we calculate # v^T tanh(W_h h_i + W_s s_t + b_attn) # where h_i is an encoder state, and s_t a decoder state. # attn_vec_size is the length of the vectors v, b_attn, (W_h h_i) and (W_s s_t). # We set it to be equal to the size of the encoder states. attention_vec_size = attn_size # Get the weight matrix W_h and apply it to each encoder state to get (W_h h_i), the encoder features if _hps.matrix_attention: w_attn = variable_scope.get_variable("w_attn", [attention_vec_size, attention_vec_size]) if _hps.intradecoder: w_dec_attn = variable_scope.get_variable("w_dec_attn", [decoder_attn_size, decoder_attn_size]) else: W_h = variable_scope.get_variable("W_h", [1, 1, attn_size, attention_vec_size]) v = variable_scope.get_variable("v", [attention_vec_size]) encoder_features = nn_ops.conv2d(_enc_states, W_h, [1, 1, 1, 1], "SAME") # shape (batch_size,max_enc_steps,1,attention_vec_size) if _hps.intradecoder: W_h_d = variable_scope.get_variable("W_h_d", [1, 1, decoder_attn_size, decoder_attn_size]) v_d = variable_scope.get_variable("v_d", [decoder_attn_size]) # Get the weight vectors v and w_c (w_c is for coverage) if use_coverage: with variable_scope.variable_scope("coverage"): w_c = variable_scope.get_variable("w_c", [1, 1, 1, attention_vec_size]) if prev_coverage is not None: # for beam search mode with coverage # reshape from (batch_size, max_enc_steps) to (batch_size, max_enc_steps, 1, 1) prev_coverage = tf.expand_dims(tf.expand_dims(prev_coverage,2),3) def attention(decoder_state, temporal_e, coverage=None): """Calculate the context vector and attention distribution from the decoder state. Args: decoder_state: state of the decoder temporal_e: store previous attentions for temporal attention mechanism coverage: Optional. Previous timestep's coverage vector, shape (batch_size, max_enc_steps, 1, 1). Returns: context_vector: weighted sum of _enc_states attn_dist: attention distribution coverage: new coverage vector. shape (batch_size, max_enc_steps, 1, 1) masked_e: store the attention score for temporal attention mechanism. """ with variable_scope.variable_scope("Attention"): # Pass the decoder state through a linear layer (this is W_s s_t + b_attn in the paper) decoder_features = linear(decoder_state, attention_vec_size, True) # shape (batch_size, attention_vec_size) decoder_features = tf.expand_dims(tf.expand_dims(decoder_features, 1), 1) # reshape to (batch_size, 1, 1, attention_vec_size) # We can't have coverage with matrix attention if not _hps.matrix_attention and use_coverage and coverage is not None: # non-first step of coverage # Multiply coverage vector by w_c to get coverage_features. coverage_features = nn_ops.conv2d(coverage, w_c, [1, 1, 1, 1], "SAME") # c has shape (batch_size, max_enc_steps, 1, attention_vec_size) # Calculate v^T tanh(W_h h_i + W_s s_t + w_c c_i^t + b_attn) e_not_masked = math_ops.reduce_sum(v * math_ops.tanh(encoder_features + decoder_features + coverage_features), [2, 3]) # shape (batch_size,max_enc_steps) masked_e = nn_ops.softmax(e_not_masked) * enc_padding_mask # (batch_size, max_enc_steps) masked_sums = tf.reduce_sum(masked_e, axis=1) # shape (batch_size) masked_e = masked_e / tf.reshape(masked_sums, [-1, 1]) # Equation 3 in if _hps.use_temporal_attention: try: len_temporal_e = temporal_e.get_shape()[0] except: len_temporal_e = 0 if len_temporal_e==0: attn_dist = masked_e else: masked_sums = tf.reduce_sum(temporal_e,axis=0)+1e-10 # if it's zero due to masking we set it to a small value attn_dist = masked_e / masked_sums # (batch_size, max_enc_steps) else: attn_dist = masked_e masked_attn_sums = tf.reduce_sum(attn_dist, axis=1) attn_dist = attn_dist / tf.reshape(masked_attn_sums, [-1, 1]) # re-normalize # Update coverage vector coverage += array_ops.reshape(attn_dist, [batch_size, -1, 1, 1]) else: if _hps.matrix_attention: # Calculate h_d * W_attn * h_i, equation 2 in https://arxiv.org/pdf/1705.04304.pdf _dec_attn = tf.unstack(tf.matmul(tf.squeeze(decoder_features,axis=[1,2]),w_attn),axis=0) # batch_size * (attention_vec_size) _enc_states_lst = tf.unstack(tf.squeeze(_enc_states,axis=2),axis=0) # batch_size * (max_enc_steps, attention_vec_size) e_not_masked = tf.squeeze(tf.stack([tf.matmul(tf.reshape(_dec,[1,-1]), tf.transpose(_enc)) for _dec, _enc in zip(_dec_attn,_enc_states_lst)]),axis=1) # (batch_size, max_enc_steps) masked_e = tf.exp(e_not_masked * enc_padding_mask) # (batch_size, max_enc_steps) else: # Calculate v^T tanh(W_h h_i + W_s s_t + b_attn) e_not_masked = math_ops.reduce_sum(v * math_ops.tanh(encoder_features + decoder_features), [2, 3]) # calculate e, (batch_size, max_enc_steps) masked_e = nn_ops.softmax(e_not_masked) * enc_padding_mask # (batch_size, max_enc_steps) masked_sums = tf.reduce_sum(masked_e, axis=1) # shape (batch_size) masked_e = masked_e / tf.reshape(masked_sums, [-1, 1]) if _hps.use_temporal_attention: try: len_temporal_e = temporal_e.get_shape()[0] except: len_temporal_e = 0 if len_temporal_e==0: attn_dist = masked_e else: masked_sums = tf.reduce_sum(temporal_e,axis=0)+1e-10 # if it's zero due to masking we set it to a small value attn_dist = masked_e / masked_sums # (batch_size, max_enc_steps) else: attn_dist = masked_e # Calculate attention distribution masked_attn_sums = tf.reduce_sum(attn_dist, axis=1) attn_dist = attn_dist / tf.reshape(masked_attn_sums, [-1, 1]) # re-normalize if use_coverage: # first step of training coverage = tf.expand_dims(tf.expand_dims(attn_dist,2),2) # initialize coverage # Calculate the context vector from attn_dist and _enc_states context_vector = math_ops.reduce_sum(array_ops.reshape(attn_dist, [batch_size, -1, 1, 1]) * _enc_states, [1, 2]) # shape (batch_size, attn_size). context_vector = array_ops.reshape(context_vector, [-1, attn_size]) return context_vector, attn_dist, coverage, masked_e def intra_decoder_attention(decoder_state, outputs): """Calculate the context vector and attention distribution from the decoder state. Args: decoder_state: state of the decoder outputs: list of decoder states for implementing intra-decoder mechanism, len(decoder_states) * (batch_size, hidden_dim) Returns: context_decoder_vector: weighted sum of _dec_states decoder_attn_dist: intra-decoder attention distribution """ attention_dec_vec_size = attn_dec_size = decoder_state.c.get_shape()[1] # hidden_dim try: len_dec_states = outputs.get_shape()[0] except: len_dec_states = 0 attention_dec_vec_size = attn_dec_size = decoder_state.c.get_shape()[1] # hidden_dim _decoder_states = tf.expand_dims(tf.reshape(outputs,[batch_size,-1,attn_dec_size]), axis=2) # now is shape (batch_size,len(decoder_states), 1, attn_size) _prev_decoder_features = nn_ops.conv2d(_decoder_states, W_h_d, [1, 1, 1, 1], "SAME") # shape (batch_size,len(decoder_states),1,attention_vec_size) with variable_scope.variable_scope("DecoderAttention"): # Pass the decoder state through a linear layer (this is W_s s_t + b_attn in the paper) try: decoder_features = linear(decoder_state, attention_dec_vec_size, True) # shape (batch_size, attention_vec_size) decoder_features = tf.expand_dims(tf.expand_dims(decoder_features, 1), 1) # reshape to (batch_size, 1, 1, attention_dec_vec_size) # Calculate v^T tanh(W_h h_i + W_s s_t + b_attn) if _hps.matrix_attention: # Calculate h_d * W_attn * h_d, equation 6 in https://arxiv.org/pdf/1705.04304.pdf _dec_attn = tf.matmul(tf.squeeze(decoder_features),w_dec_attn) # (batch_size, decoder_attn_size) _dec_states_lst = tf.unstack(tf.reshape(_prev_decoder_features,[batch_size,-1,decoder_attn_size])) # batch_size * (len(decoder_states), decoder_attn_size) e_not_masked = tf.reshape(tf.stack([tf.matmul(_dec_attn, tf.transpose(k)) for k in _dec_states_lst]),[batch_size,-1]) # (batch_size, len(decoder_states)) masked_e = tf.exp(e_not_masked * dec_padding_mask[:,:len_dec_states]) # (batch_size, len(decoder_states)) else: # Calculate v^T tanh(W_h h_i + W_s s_t + b_attn) e_not_masked = math_ops.reduce_sum(v_d * math_ops.tanh(_prev_decoder_features + decoder_features), [2, 3]) # calculate e, (batch_size,len(decoder_states)) masked_e = nn_ops.softmax(e_not_masked) * dec_padding_mask[:,:len_dec_states] # (batch_size,len(decoder_states)) if len_dec_states <= 1: masked_e = array_ops.ones([batch_size,1]) # first step is filled with equal values masked_sums = tf.reshape(tf.reduce_sum(masked_e,axis=1),[-1,1]) # (batch_size,1), # if it's zero due to masking we set it to a small value decoder_attn_dist = masked_e / masked_sums # (batch_size,len(decoder_states)) context_decoder_vector = math_ops.reduce_sum(array_ops.reshape(decoder_attn_dist, [batch_size, -1, 1, 1]) * _decoder_states, [1, 2]) # (batch_size, attn_size) context_decoder_vector = array_ops.reshape(context_decoder_vector, [-1, attn_dec_size]) # (batch_size, attn_size) except: return array_ops.zeros([batch_size, decoder_attn_size]), array_ops.zeros([batch_size, 0]) return context_decoder_vector, decoder_attn_dist outputs = [] temporal_e = [] attn_dists = [] vocab_scores = [] vocab_dists = [] final_dists = [] p_gens = [] samples = [] # this holds the words chosen by sampling based on the final distribution for each decoding step, list of max_dec_steps of (batch_size, 1) greedy_search_samples = [] # this holds the words chosen by greedy search (taking the max) on the final distribution for each decoding step, list of max_dec_steps of (batch_size, 1) sampling_rewards = [] # list of size max_dec_steps (batch_size, k) greedy_rewards = [] # list of size max_dec_steps (batch_size, k) state = _dec_in_state coverage = prev_coverage # initialize coverage to None or whatever was passed in context_vector = array_ops.zeros([batch_size, attn_size]) context_decoder_vector = array_ops.zeros([batch_size, decoder_attn_size]) context_vector.set_shape([None, attn_size]) # Ensure the second shape of attention vectors is set. if initial_state_attention: # true in decode mode # Re-calculate the context vector from the previous step so that we can pass it through a linear layer with this step's input to get a modified version of the input context_vector, _, coverage, _ = attention(_dec_in_state, tf.stack(prev_encoder_es,axis=0), coverage) # in decode mode, this is what updates the coverage vector if _hps.intradecoder: context_decoder_vector, _ = intra_decoder_attention(_dec_in_state, tf.stack(prev_decoder_outputs,axis=0)) for i, inp in enumerate(emb_dec_inputs): tf.logging.info("Adding attention_decoder timestep %i of %i", i, len(emb_dec_inputs)) if i > 0: variable_scope.get_variable_scope().reuse_variables() if _hps.mode in ['train','eval'] and _hps.scheduled_sampling and i > 0: # start scheduled sampling after we received the first decoder's output # modify the input to next decoder using scheduled sampling if FLAGS.scheduled_sampling_final_dist: inp = scheduled_sampling(_hps, sampling_probability, final_dist, embedding, inp, alpha) else: inp = scheduled_sampling_vocab_dist(_hps, sampling_probability, vocab_dist, embedding, inp, alpha) # Merge input and previous attentions into one vector x of the same size as inp emb_dim = inp.get_shape().with_rank(2)[1] if emb_dim is None: raise ValueError("Could not infer input size from input: %s" % inp.name) x = linear([inp] + [context_vector], emb_dim, True) # Run the decoder RNN cell. cell_output = decoder state cell_output, state = cell(x, state) # Run the attention mechanism. if i == 0 and initial_state_attention: # always true in decode mode with variable_scope.variable_scope(variable_scope.get_variable_scope(), reuse=True): # you need this because you've already run the initial attention(...) call context_vector, attn_dist, _, masked_e = attention(state, tf.stack(prev_encoder_es,axis=0), coverage) # don't allow coverage to update if _hps.intradecoder: context_decoder_vector, _ = intra_decoder_attention(state, tf.stack(prev_decoder_outputs,axis=0)) else: context_vector, attn_dist, coverage, masked_e = attention(state, tf.stack(temporal_e,axis=0), coverage) if _hps.intradecoder: context_decoder_vector, _ = intra_decoder_attention(state, tf.stack(outputs,axis=0)) attn_dists.append(attn_dist) temporal_e.append(masked_e) with variable_scope.variable_scope("combined_context"): if _hps.intradecoder: context_vector = linear([context_vector] + [context_decoder_vector], attn_size, False) # Calculate p_gen if pointer_gen: with tf.variable_scope('calculate_pgen'): p_gen = linear([context_vector, state.c, state.h, x], 1, True) # Tensor shape (batch_size, 1) p_gen = tf.sigmoid(p_gen) p_gens.append(p_gen) # Concatenate the cell_output (= decoder state) and the context vector, and pass them through a linear layer # This is V[s_t, h*_t] + b in the paper with variable_scope.variable_scope("AttnOutputProjection"): output = linear([cell_output] + [context_vector], cell.output_size, True) outputs.append(output) # Add the output projection to obtain the vocabulary distribution with tf.variable_scope('output_projection'): if i > 0: tf.get_variable_scope().reuse_variables() trunc_norm_init = tf.truncated_normal_initializer(stddev=_hps.trunc_norm_init_std) w_out = tf.get_variable('w', [_hps.dec_hidden_dim, v_size], dtype=tf.float32, initializer=trunc_norm_init) #w_t_out = tf.transpose(w) v_out = tf.get_variable('v', [v_size], dtype=tf.float32, initializer=trunc_norm_init) if i > 0: tf.get_variable_scope().reuse_variables() if FLAGS.share_decoder_weights: # Eq. 13 in https://arxiv.org/pdf/1705.04304.pdf w_out = tf.transpose( math_ops.tanh(linear([embedding] + [tf.transpose(w_out)], _hps.dec_hidden_dim, bias=False))) score = tf.nn.xw_plus_b(output, w_out, v_out) if _hps.scheduled_sampling and not _hps.greedy_scheduled_sampling: # Gumbel reparametrization trick: https://arxiv.org/abs/1704.06970 U = tf.random_uniform(score.get_shape(),10e-12,(1-10e-12)) # add a small number to avoid log(0) G = -tf.log(-tf.log(U)) score = score + G vocab_scores.append(score) # apply the linear layer vocab_dist = tf.nn.softmax(score) vocab_dists.append(vocab_dist) # The vocabulary distributions. List length max_dec_steps of (batch_size, vsize) arrays. The words are in the order they appear in the vocabulary file. # For pointer-generator model, calc final distribution from copy distribution and vocabulary distribution if _hps.pointer_gen: final_dist = _calc_final_dist(_hps, v_size, _max_art_oovs, _enc_batch_extend_vocab, p_gen, vocab_dist, attn_dist) else: # final distribution is just vocabulary distribution final_dist = vocab_dist final_dists.append(final_dist) # get the sampled token and greedy token # this will take the final_dist and sample from it for a total count of k (k samples) one_hot_k_samples = tf.distributions.Multinomial(total_count=1., probs=final_dist).sample( _hps.k) # sample k times according to https://arxiv.org/pdf/1705.04304.pdf, size (k, batch_size, extended_vsize) k_argmax = tf.argmax(one_hot_k_samples, axis=2, output_type=tf.int32) # (k, batch_size) k_sample = tf.transpose(k_argmax) # shape (batch_size, k) greedy_search_prob, greedy_search_sample = tf.nn.top_k(final_dist, k=_hps.k) # (batch_size, k) greedy_search_samples.append(greedy_search_sample) samples.append(k_sample) if FLAGS.use_discounted_rewards: _sampling_rewards = [] _greedy_rewards = [] for _ in range(_hps.k): rl_fscore = tf.reshape(rouge_l_fscore(tf.transpose(tf.stack(samples)[:, :, _]), target_batch), [-1, 1]) # shape (batch_size, 1) _sampling_rewards.append(tf.reshape(rl_fscore, [-1, 1])) rl_fscore = tf.reshape(rouge_l_fscore(tf.transpose(tf.stack(greedy_search_samples)[:, :, _]), target_batch), [-1, 1]) # shape (batch_size, 1) _greedy_rewards.append(tf.reshape(rl_fscore, [-1, 1])) sampling_rewards.append(tf.squeeze(tf.stack(_sampling_rewards, axis=1), axis = -1)) # (batch_size, k) greedy_rewards.append(tf.squeeze(tf.stack(_greedy_rewards, axis=1), axis = -1)) # (batch_size, k) if FLAGS.use_discounted_rewards: sampling_rewards = tf.stack(sampling_rewards) greedy_rewards = tf.stack(greedy_rewards) else: _sampling_rewards = [] _greedy_rewards = [] for _ in range(_hps.k): rl_fscore = rouge_l_fscore(tf.transpose(tf.stack(samples)[:, :, _]), target_batch) # shape (batch_size, 1) _sampling_rewards.append(tf.reshape(rl_fscore, [-1, 1])) rl_fscore = rouge_l_fscore(tf.transpose(tf.stack(greedy_search_samples)[:, :, _]), target_batch) # shape (batch_size, 1) _greedy_rewards.append(tf.reshape(rl_fscore, [-1, 1])) sampling_rewards = tf.squeeze(tf.stack(_sampling_rewards, axis=1), axis=-1) # (batch_size, k) greedy_rewards = tf.squeeze(tf.stack(_greedy_rewards, axis=1), axis=-1) # (batch_size, k) # If using coverage, reshape it if coverage is not None: coverage = array_ops.reshape(coverage, [batch_size, -1]) return ( outputs, state, attn_dists, p_gens, coverage, vocab_scores, final_dists, samples, greedy_search_samples, temporal_e, sampling_rewards, greedy_rewards)
def _add_shared_loss_op(self): # Calculate the loss with tf.variable_scope('shared_loss'): # Calculate the loss per step # This is fiddly; we use tf.gather_nd to pick out the probabilities of the gold target words #### added by [email protected]: we just calculate these to monitor pgen_loss throughout time loss_per_step = [] # will be list length max_dec_steps containing shape (batch_size) batch_nums = tf.range(0, limit=self._hps.batch_size) # shape (batch_size) for dec_step, dist in enumerate(self.final_dists): targets = self._target_batch[:,dec_step] # The indices of the target words. shape (batch_size) indices = tf.stack( (batch_nums, targets), axis=1) # shape (batch_size, 2) gold_probs = tf.gather_nd(dist, indices) # shape (batch_size). prob of correct words on this step losses = -tf.log(gold_probs) loss_per_step.append(losses) self._pgen_loss = _mask_and_avg(loss_per_step, self._dec_padding_mask) self.variable_summaries('pgen_loss', self._pgen_loss) # Adding Q-Estimation to CE loss in Actor-Critic Model if self._hps.ac_training: # Calculating Actor-Critic loss # Here, we multiple the Q-estimation for each token to its respective probability loss_per_step = [] # will be list length k each containing a list of shape <=max_dec_steps which each has the shape (batch_size) q_loss_per_step = [] # will be list length k each containing a list of shape <=max_dec_steps which each has the shape (batch_size) batch_nums = tf.range(0, limit=self._hps.batch_size) # shape (batch_size) unstacked_q = tf.unstack(self._q_estimates, axis =1) # list of k with size (batch_size, <=max_dec_steps, vsize_extended) for sample_id in range(self._hps.k): loss_per_sample = [] # length <=max_dec_steps of batch_sizes q_loss_per_sample = [] # length <=max_dec_steps of batch_sizes q_val_per_sample = tf.unstack(unstacked_q[sample_id], axis =1) # list of <=max_dec_step (batch_size, vsize_extended) for dec_step, (dist, q_value) in enumerate(zip(self.final_dists, q_val_per_sample)): targets = tf.squeeze(self.samples[dec_step][:,sample_id]) # The indices of the sampled words. shape (batch_size) indices = tf.stack((batch_nums, targets), axis=1) # shape (batch_size, 2) gold_probs = tf.gather_nd(dist, indices) # shape (batch_size). prob of correct words on this step losses = -tf.log(gold_probs) dist_q_val = -tf.log(dist) * q_value q_losses = tf.gather_nd(dist_q_val, indices) # shape (batch_size). prob of correct words on this step loss_per_sample.append(losses) q_loss_per_sample.append(q_losses) loss_per_step.append(loss_per_sample) q_loss_per_step.append(q_loss_per_sample) with tf.variable_scope('reinforce_loss'): #### this is the actual loss self._rl_avg_logprobs = tf.reduce_mean([_mask_and_avg(loss_per_sample, self._dec_padding_mask) for loss_per_sample in loss_per_step]) self._rl_loss = tf.reduce_mean([_mask_and_avg(q_loss_per_sample, self._dec_padding_mask) for q_loss_per_sample in q_loss_per_step]) # Eq. 34 in https://arxiv.org/pdf/1805.09461.pdf self._reinforce_shared_loss = self._eta * self._rl_loss + (tf.constant(1.,dtype=tf.float32) - self._eta) * self._pgen_loss # equation 16 in https://arxiv.org/pdf/1705.04304.pdf #### the following is only for monitoring purposes self.variable_summaries('reinforce_avg_logprobs', self._rl_avg_logprobs) self.variable_summaries('reinforce_loss', self._rl_loss) self.variable_summaries('reinforce_shared_loss', self._reinforce_shared_loss) # Adding Self-Critic Reward to CE loss in Policy-Gradient Model if self._hps.rl_training: #### Calculating the reinforce loss according to Eq. 15 in https://arxiv.org/pdf/1705.04304.pdf loss_per_step = [] # will be list length max_dec_steps*k containing shape (batch_size) rl_loss_per_step = [] # will be list length max_dec_steps*k containing shape (batch_size) batch_nums = tf.range(0, limit=self._hps.batch_size) # shape (batch_size) self._sampled_rouges = [] self._greedy_rouges = [] self._reward_diff = [] for _ in range(self._hps.k): self._sampled_rouges.append(rouge_l_fscore(self.sampled_sentences[:, _, :], self._target_batch)) self._greedy_rouges.append(rouge_l_fscore(self.greedy_search_sentences[:, _, :], self._target_batch)) self._reward_diff.append(self._sampled_rouges[_] - self._greedy_rouges[_]) for dec_step, dist in enumerate(self.final_dists): _targets = self.samples[dec_step] # The indices of the sampled words. shape (batch_size, k) for _k, targets in enumerate(tf.unstack(_targets,axis=1)): # list of k samples of size (batch_size) indices = tf.stack( (batch_nums, targets), axis=1) # shape (batch_size, 2) gold_probs = tf.gather_nd(dist, indices) # shape (batch_size). prob of correct words on this step losses = -tf.log(gold_probs) loss_per_step.append(losses) # Equation 15 in https://arxiv.org/pdf/1705.04304.pdf # Equal reward for all tokens rl_losses = -tf.log(gold_probs) * self._reward_diff[_k] #rl_losses = -tf.log(gold_probs) * (self._sampled_sentence_r_values[_k]-self._greedy_sentence_r_values[_k]) rl_loss_per_step.append(rl_losses) # new size: (k, max_dec_steps, batch_size) rl_loss_per_step = tf.unstack( tf.transpose(tf.reshape(rl_loss_per_step, [-1, self._hps.k, self._hps.batch_size]),perm=[1,0,2])) loss_per_step = tf.unstack( tf.transpose(tf.reshape(loss_per_step, [-1, self._hps.k, self._hps.batch_size]), perm=[1, 0, 2])) with tf.variable_scope('reinforce_loss'): self._rl_avg_logprobs = [] self._rl_loss = [] for _k in range(self._hps.k): self._rl_avg_logprobs.append(_mask_and_avg(tf.unstack(loss_per_step[_k]), self._dec_padding_mask)) self._rl_loss.append(_mask_and_avg(tf.unstack(tf.reshape(rl_loss_per_step[_k], [self._hps.max_dec_steps, self._hps.batch_size])), self._dec_padding_mask)) self._rl_avg_logprobs = tf.reduce_mean(self._rl_avg_logprobs) self._rl_loss = tf.reduce_mean(self._rl_loss) # We multiply the ROUGE difference of sampling vs greedy sentence to the loss of all tokens in the sequence # Eq. 16 in https://arxiv.org/pdf/1705.04304.pdf and Eq. 34 in https://arxiv.org/pdf/1805.09461.pdf self._reinforce_shared_loss = self._eta * self._rl_loss + (tf.constant(1.,dtype=tf.float32) - self._eta) * self._pgen_loss #### the following is only for monitoring purposes self.variable_summaries('reinforce_avg_logprobs', self._rl_avg_logprobs) self.variable_summaries('reinforce_loss', self._rl_loss) self.variable_summaries('reinforce_sampled_r_value', tf.reduce_mean(self._sampled_rouges)) self.variable_summaries('reinforce_greedy_r_value', tf.reduce_mean(self._greedy_rouges)) self.variable_summaries('reinforce_r_diff', tf.reduce_mean(self._reward_diff)) self.variable_summaries('reinforce_shared_loss', self._reinforce_shared_loss) # Calculate coverage loss from the attention distributions if self._hps.coverage: with tf.variable_scope('coverage_loss'): self._coverage_loss = _coverage_loss(self.attn_dists, self._dec_padding_mask) self.variable_summaries('coverage_loss', self._coverage_loss) if self._hps.rl_training or self._hps.ac_training: with tf.variable_scope('reinforce_loss'): self._reinforce_cov_total_loss = self._reinforce_shared_loss + self._hps.cov_loss_wt * self._coverage_loss self.variable_summaries('reinforce_coverage_loss', self._reinforce_cov_total_loss) if self._hps.pointer_gen: self._pointer_cov_total_loss = self._pgen_loss + self._hps.cov_loss_wt * self._coverage_loss self.variable_summaries('pointer_coverage_loss', self._pointer_cov_total_loss)