def build_seq2seq(input_seqs, target_seqs, filtered_target_seqs, input_seq_lengths, target_seq_lengths, is_training): with tf.variable_scope("seq2seq"): with tf.device('/cpu:0'): reuse = False if get_PPL: keep_prob = tf.convert_to_tensor(1.0) else: keep_prob = get_keep_prob(dropout_rate, is_training) sequence_mask = get_sequence_mask(target_seq_lengths) unk_mask = get_mask(target_seqs, unk_indices) decoder_mask = tf.logical_and(sequence_mask, tf.logical_not(unk_mask)) decoder_mask_float = tf.cast(decoder_mask, tf.float32) # Embed inputs with tf.variable_scope("embedding"): embedding = create_embedding(embedding_word2vec_politeness, embedding_word2vec_movie, shared_vocab_size_politeness, shared_vocab_size_movie, new_vocab_size_politeness, new_vocab_size_movie, "seq2seq") embedded_input_seqs = tf.nn.embedding_lookup( embedding, input_seqs) embedded_target_seqs = tf.nn.embedding_lookup( embedding, target_seqs) # Optimizer optimizer = tf.train.AdamOptimizer(learning_rate) tower_grads = [] if credit_assignment: tower_grads_polite = [] sample_ids_lst = [] final_lengths_lst = [] sampled_sample_ids_lst = [] sampled_final_lengths_lst = [] reuse = False trainable_variables = [] num_tokens_lst = [] total_losses = [] for i in xrange(num_gpus): with tf.device("/gpu:%d" % (gpu_start_index + i)): with tf.variable_scope("seq2seq"): if (i == 1): reuse = True start = i * batch_size_per_gpu end = start + batch_size_per_gpu input_max_seq_length = tf.reduce_max( input_seq_lengths[start:end]) target_max_seq_length = tf.reduce_max( target_seq_lengths[start:end]) with tf.variable_scope("encoder", reuse=reuse): cell_fw = create_MultiRNNCell([hidden_size_encoder] * (num_layers_encoder // 2), keep_prob, num_proj=None, reuse=reuse) cell_bw = create_MultiRNNCell([hidden_size_encoder] * (num_layers_encoder // 2), keep_prob, num_proj=None, reuse=reuse) (encoder_outputs_original, encoder_final_state_original ) = bidirecitonal_dynamic_lstm( cell_fw, cell_bw, embedded_input_seqs[ start:end, :input_max_seq_length, :], input_seq_lengths[start:end]) [ encoder_outputs, encoder_seq_lengths, encoder_final_state ] = tf.cond(is_training, lambda: [ encoder_outputs_original, input_seq_lengths[start:end], encoder_final_state_original ], lambda: [ tf.contrib.seq2seq.tile_batch(encoder_outputs_original, beam_width), tf.contrib.seq2seq.tile_batch( input_seq_lengths[start:end], beam_width), tile_multi_cell_state(encoder_final_state_original) ]) # only works for decoder that has >1 layers! with tf.variable_scope("decoder", reuse=reuse): decoder_cell = create_MultiRNNCell( [hidden_size_decoder] * (num_layers_decoder), keep_prob, num_proj=vocab_size, memory=encoder_outputs, memory_seq_lengths=encoder_seq_lengths, reuse=reuse) decoder_zero_state = tf.cond( is_training, lambda: decoder_cell.zero_state( batch_size_per_gpu, tf.float32), lambda: decoder_cell.zero_state( batch_size_per_gpu * beam_width, tf.float32)) state_last = decoder_zero_state[-1].clone( cell_state=encoder_final_state[-1]) state_previous = encoder_final_state[:-1] decoder_initial_state = state_previous + ( state_last, ) # concat tuples # training helper (for teacher forcing) helper_train = tf.contrib.seq2seq.TrainingHelper( embedded_target_seqs[ start:end, :target_max_seq_length - 1, :], # get rid of end_token target_seq_lengths[start:end] - 1) # the length is thus decreased by 1 (decoder_outputs_train, _) = decode(decoder_cell, helper_train, initial_state=decoder_initial_state) (logits, _) = decoder_outputs_train # Get trainable_variables # (up to now we already have all the seq2seq trainable vars) if trainable_variables == []: trainable_variables = tf.get_collection( tf.GraphKeys.TRAINABLE_VARIABLES, scope="seq2seq") loss_ML = tf.contrib.seq2seq.sequence_loss( logits, target_seqs[ start:end, 1:target_max_seq_length], # get rid of start_token decoder_mask_float[start:end, 1:target_max_seq_length]) num_tokens = tf.reduce_sum( decoder_mask_float[start:end, 1:target_max_seq_length]) num_tokens_lst.append(num_tokens) total_loss = loss_ML * num_tokens total_losses.append(total_loss) if polite_training: helper_sample = tf.contrib.seq2seq.SampleEmbeddingHelper( embedding, start_tokens[start:end], end_token) (decoder_outputs_sample, final_lengths_sample) = decode( decoder_cell, helper_sample, decoder_initial_state) (logits_sample, sample_ids_sample) = decoder_outputs_sample max_final_lengths_sample = tf.reduce_max( final_lengths_sample) sampled_sample_ids_lst.append( pad_and_truncate(sample_ids_sample, final_lengths_sample)) sampled_final_lengths_lst.append(final_lengths_sample) # Compute sampled sequence loss WITHOUT averaging (will do that later) decoder_mask_sample = get_sequence_mask( final_lengths_sample, dtype=tf.float32) seq_losses_sample = tf.contrib.seq2seq.sequence_loss( logits_sample, sample_ids_sample, decoder_mask_sample, average_across_timesteps=False, average_across_batch=False) if polite_training: with tf.variable_scope( "classifier"): # jump back to the classifier scope # Filter out tokens that the classifier doesn't know vocab_mask = tf.cast( sample_ids_sample < vocab_size_politeness, tf.int32) sample_ids_sample_classifier = sample_ids_sample * vocab_mask # Feed sampled ids to classifier (scores_RL, credit_weights_RL) = build_classifier( sample_ids_sample_classifier, final_lengths_sample, reuse) # Stop gradients from propagating back scores_RL_stop = tf.stop_gradient(scores_RL) credit_weights_RL_stop = tf.stop_gradient( credit_weights_RL) if thresholding: # Filter scores that are >= threshold and <= 1 - threshold filtered_scores_RL = tf.map_fn(filter_with_threshold, scores_RL_stop) else: filtered_scores_RL = scores_RL_stop with tf.variable_scope("seq2seq"): with tf.variable_scope("decoder", reuse=reuse): # Get valid mask for sampled sequence decoder_mask_classifier = tf.cast( tf.not_equal(sample_ids_sample, 0), tf.float32 ) # propagate back the whole sentence (including <end>) tiled_scores = tf.tile( # tile scores to 2D tf.expand_dims(filtered_scores_RL - baseline, axis=1), [1, max_final_lengths_sample]) if flip_polite: # if we actually want a rude dialogue system tiled_scores = -1.0 * tiled_scores # Compute seq losses for polite-RL seq_losses_classifier = ( beta * seq_losses_sample * decoder_mask_classifier / tf.reduce_sum(decoder_mask_classifier) * tiled_scores) if credit_assignment: grads_polite = tf.gradients( seq_losses_classifier, trainable_variables, grad_ys=credit_weights_RL_stop ) # credit weights as initial gradients grads_polite = zip_lsts( [grads_polite, trainable_variables]) tower_grads_polite.append(grads_polite) else: loss_polite = tf.reduce_sum(seq_losses_classifier) else: credit_weights_RL_stop = None with tf.variable_scope("seq2seq"): with tf.variable_scope("decoder", reuse=reuse): # Infer branch (beam search!) beam_search_decoder = tf.contrib.seq2seq.BeamSearchDecoder( decoder_cell, embedding, start_tokens[start:end], end_token, decoder_initial_state, beam_width, length_penalty_weight=length_penalty_weight) output_beam = tf.contrib.seq2seq.dynamic_decode( beam_search_decoder, # impute_finished=True, # cannot be used with Beamsearch maximum_iterations=max_iterations, swap_memory=True) sample_ids = output_beam[0].predicted_ids[:, :, 0] final_lengths = output_beam[2][:, 0] sample_ids_lst.append( pad_and_truncate(sample_ids, final_lengths)) final_lengths_lst.append(final_lengths) with tf.device("/gpu:%d" % (gpu_start_index + i)): with tf.variable_scope("seq2seq", reuse=reuse): # Compute loss loss = loss_ML if polite_training and not credit_assignment: loss = loss + loss_polite # Compute tower gradients grads = compute_grads(loss, optimizer, trainable_variables) tower_grads.append(grads) with tf.device('/cpu:0'): with tf.variable_scope("seq2seq"): # Concat sample ids and their respective lengths batch_sample_ids = tf.concat(sample_ids_lst, axis=0) batch_final_lengths = tf.concat(final_lengths_lst, axis=0) if polite_training: batch_sampled_sample_ids = tf.concat(sampled_sample_ids_lst, axis=0) batch_total_loss = tf.add_n(total_losses) batch_num_tokens = tf.add_n(num_tokens_lst) # Thus, the effective batch size is actually batch_size_per_gpu if polite_training and credit_assignment: apply_gradients_op = apply_multiple_grads( optimizer, [tower_grads, tower_grads_polite]) else: apply_gradients_op = apply_grads(optimizer, tower_grads) return (batch_sample_ids, batch_final_lengths, batch_total_loss, batch_num_tokens, apply_gradients_op, credit_weights_RL_stop, embedding)
def __init__( self, batch_size, vocab_size, embedding_size, hidden_size_encoder, hidden_size_decoder, max_iterations, start_token, end_token, unk_indices, num_layers_encoder=1, num_layers_decoder=1, attention_size=512, attention_layer_size=256, beam_width=10, length_penalty_weight=1.0, gpu_start_index=0, num_gpus=1, # set to 1 when testing learning_rate=0.001, clipping_threshold=5.0, feed_both_examples=False, use_max_margin=False, max_margin_weight=1.0, margin=0.1, reward_clipping_threshold=1.0, backward=False, # whether we are training a backward model feed_tensors=[], # when provided, placeholders are not used use_MMI_reward=False, MMI_weight=0.00, use_reranking_reward=False, reranking_weight=0.00, num_samples_reranking=2, use_gleu_reward=False, gleu_weight=0.00, softmax_temperature=1.0, beam_search=True ): # how many samples to use for baseline (RL training) self.batch_size = batch_size self.vocab_size = vocab_size self.embedding_size = embedding_size self.hidden_size_encoder = hidden_size_encoder self.hidden_size_decoder = hidden_size_decoder assert self.hidden_size_encoder * 2 == self.hidden_size_decoder self.max_iterations = max_iterations self.start_token = start_token self.end_token = end_token self.unk_indices = unk_indices self.num_layers_encoder = num_layers_encoder self.num_layers_decoder = num_layers_decoder self.attention_size = attention_size self.attention_layer_size = attention_layer_size self.beam_width = beam_width self.length_penalty_weight = length_penalty_weight self.gpu_start_index = gpu_start_index self.num_gpus = num_gpus self.learning_rate = learning_rate self.clipping_threshold = clipping_threshold self.feed_both_examples = feed_both_examples if self.feed_both_examples: assert self.batch_size % 2 == 0 self.half_batch_size = self.batch_size // 2 self.use_max_margin = use_max_margin if self.use_max_margin: assert self.feed_both_examples # if Should-Change, then feed_both_examples must be True self.max_margin_weight = max_margin_weight self.margin = margin self.beam_search = beam_search assert self.batch_size % self.num_gpus == 0 self.batch_size_per_gpu = self.batch_size // self.num_gpus self.feed_tensors = feed_tensors self.use_MMI_reward = use_MMI_reward self.MMI_weight = MMI_weight self.use_reranking_reward = use_reranking_reward self.reranking_weight = reranking_weight self.num_samples_reranking = num_samples_reranking self.use_gleu_reward = use_gleu_reward self.gleu_weight = gleu_weight self.RL_training = self.use_MMI_reward or self.use_reranking_reward or self.use_gleu_reward self.softmax_temperature = softmax_temperature if self.feed_both_examples: print("Feeding both examples...") assert self.batch_size % 2 == 0 self.norm_batch_size = self.batch_size // 2 else: self.norm_batch_size = self.batch_size # when feeding both examples, only first half are norm inputs if self.use_max_margin: print("Max margin weight: {}, margin: {}".format( self.max_margin_weight, self.margin)) # We are only performing RL training on the norm_batch_size part, which may or may not be the whole batch if self.use_MMI_reward: print("MMI weight:", self.MMI_weight) # self.softmax_temperature = 0.5 # print("softmax_temperature changed to {}".format(self.softmax_temperature)) else: self.MMI_weight = 0.0 if self.use_reranking_reward: print("Neural Reranking reward weight:", self.reranking_weight) else: self.reranking_weight = 0.0 if self.use_gleu_reward: print("GLEU reward weight:", self.gleu_weight) else: self.gleu_weight = 0.0 self.ML_weight = 1.0 - (self.MMI_weight + self.reranking_weight + self.gleu_weight) self.trainable_variables = [] if self.use_MMI_reward: self.trainable_variables_backward = [] # For a backward model, the namespace will be "seq2seq_backward" extra_str = "_backward" if backward else "" self.main_scope = "seq2seq" + extra_str + "/" with tf.device("/gpu:%d" % self.gpu_start_index): self.create_placeholders() # Tile only if we are not training and using beam search self.tile = tf.logical_and(tf.logical_not(self.is_training), self.beam_search) self.global_step = tf.get_variable(self.main_scope + "global_step", initializer=0, dtype=tf.int32, trainable=False) # Note: if feeding both examples, the first dimension of total_loss is twice as that of # loss_RL's (self.total_loss, max_margin_loss, num_tokens, loss_MMI, loss_gleu, loss_reranking, num_tokens_RL, self.batch_sample_ids_beam, self.batch_final_lengths_beam) = self.one_iteration( self.source, self.source_length, self.target, self.target_length, self.start_tokens) # This part is for monitoring PPL, not for training. if self.feed_both_examples: self.batch_num_tokens = num_tokens / 2.0 # We montior ML losses for both norm- and adv-data self.batch_total_loss = (tf.reduce_sum( self.norm(self.total_loss)), tf.reduce_sum( self.adv(self.total_loss))) else: self.batch_num_tokens = num_tokens self.batch_total_loss = tf.reduce_sum(self.total_loss) loss_terms = [] # when using max_margin, it must be a Should-Change strategy if self.use_max_margin: loss_ML = self.norm( self.total_loss ) # in this case we don't want to train on (adv-S, T) pairs num_tokens_ML = num_tokens / 2.0 else: # if self.feed_both_examples: # # This is just for code readability # # We could have just written loss_ML = self.total_loss / 2.0 # loss_ML = (self.norm(self.total_loss) + self.adv(self.total_loss)) / 2.0 # else: # loss_ML = self.total_loss loss_ML = self.total_loss num_tokens_ML = num_tokens loss_terms.append(self.ML_weight * tf.reduce_sum(loss_ML) / num_tokens_ML) if self.use_max_margin: # need to scale max_margin_weight by ML_weight to make the training stable # (instead of scaling with loss_ML + loss_RL) loss_terms.append( self.max_margin_weight * self.ML_weight * # tf.reduce_sum(max_margin_loss) / num_tokens_ML) tf.reduce_mean(max_margin_loss)) if self.RL_training and not self.use_max_margin: num_tokens_RL = num_tokens_RL / 2.0 # effectively double the RL_loss if self.use_MMI_reward: loss_terms.append(self.MMI_weight * tf.reduce_sum(loss_MMI) / num_tokens_RL) if self.use_reranking_reward: loss_terms.append(self.reranking_weight * tf.reduce_sum(loss_reranking) / num_tokens_RL) if self.use_gleu_reward: loss_terms.append(self.gleu_weight * tf.reduce_sum(loss_gleu) / num_tokens_RL) assert loss_terms != [] loss = tf.add_n(loss_terms) optimizer = tf.train.AdamOptimizer(self.learning_rate) grads = compute_grads(loss, optimizer, self.trainable_variables) self.apply_gradients_op = apply_grads( optimizer, [grads], clipping_threshold=self.clipping_threshold, global_step=self.global_step)
def __init__(self, batch_size, vocab_size, embedding_size, hidden_size_encoder, hidden_size_context, hidden_size_decoder, max_iterations, max_dialogue_length, start_token, end_token, unk_indices, num_layers_encoder=1, num_layers_context=1, num_layers_decoder=1, attention_size=512, attention_layer_size=256, beam_search=False, beam_width=10, length_penalty_weight=1.0, gpu_start_index=0, learning_rate=0.001, clipping_threshold=5.0, feed_both_examples=False, use_max_margin=False, max_margin_weight=1.0, margin=0.5): self.batch_size = batch_size self.half_batch_size = self.batch_size // 2 self.vocab_size = vocab_size self.embedding_size = embedding_size self.hidden_size_encoder = hidden_size_encoder self.hidden_size_context = hidden_size_context self.hidden_size_decoder = hidden_size_decoder self.dim_z = self.hidden_size_context # this decision is arbitrary self.create_context_initial_state_var = functools.partial( tf.get_variable, initializer=tf.zeros([self.batch_size, self.hidden_size_context]), dtype=tf.float32, trainable=False) self.dense = functools.partial(tf.layers.dense, units=self.dim_z, use_bias=True) self.max_iterations = max_iterations self.max_dialogue_length = max_dialogue_length assert self.max_dialogue_length > 0 self.start_token = start_token self.end_token = end_token self.unk_indices = unk_indices self.num_layers_encoder = num_layers_encoder self.num_layers_context = num_layers_context self.num_layers_decoder = num_layers_decoder self.attention_size = attention_size self.attention_layer_size = attention_layer_size self.beam_search = beam_search self.beam_width = beam_width self.length_penalty_weight = length_penalty_weight self.clipping_threshold = clipping_threshold self.feed_both_examples = feed_both_examples if self.feed_both_examples: assert self.batch_size % 2 == 0 self.use_max_margin = use_max_margin if self.use_max_margin: assert self.feed_both_examples self.max_margin_weight = max_margin_weight self.margin = margin self.prior_str = "prior" self.posterior_str = "posterior" """ context_input_acc put context inputs in reverse order. • When adding inputs, just concat from left • When using, we do not need to reverse it back because we are using bidirectional-lstm Note: If we don't use bidirectional-lstm, then we will need tf.reverse_sequence() """ self.trainable_variables = [] with tf.variable_scope("seq2seq"): self.create_placeholders() # Tile only if we are not training and using beam search self.tile = tf.logical_and(tf.logical_not(self.is_training), self.beam_search) context_input_acc = tf.get_variable( "context_input_acc", initializer=tf.zeros( [self.batch_size, 2, self.hidden_size_encoder * 2], dtype=tf.float32), trainable=False) self.global_step = tf.get_variable("global_step", initializer=0, dtype=tf.int32, trainable=False) # max_num_turns better not be less than 2, since we may just lose a whole dimension (i.e., axis=1)? max_num_turns = tf.maximum(tf.reduce_max(self.start_turn_index), 2) context_input_mask = tf.tile( tf.reshape(tf.greater(self.start_turn_index, 0), [self.batch_size, 1, 1]), # expand two dims [1, max_num_turns, self.hidden_size_encoder * 2]) # This multiplication resets context input that have start_turn_index == 0 previous_context_input = context_input_acc[:, : max_num_turns, :] * tf.cast( context_input_mask, tf.float32) optimizer = tf.train.AdamOptimizer(learning_rate) with tf.device("/gpu:%d" % gpu_start_index): (kl_loss, total_loss, num_tokens, max_margin_loss, context_input, sample_ids, final_lengths) = self.one_iteration(self.dialogue, self.turn_length, previous_context_input, self.start_turn_index, self.start_tokens) kl_loss_weight = tf.cond( self.tile, lambda: tf.minimum( 1.0 / 75000.0 * tf.cast(self.global_step, tf.float32), 1.0), lambda: 1.0) if self.use_max_margin: kl_loss = kl_loss[:self.half_batch_size] weighted_kl_loss = kl_loss_weight * tf.reduce_mean(kl_loss) if self.use_max_margin: num_tokens = num_tokens / 2.0 # Note: here total_loss is a 1-D vector (already summed with axis=1) cross_ent_loss = tf.reduce_sum( total_loss[:self.half_batch_size]) adv_cross_ent_loss = tf.reduce_sum( total_loss[self.half_batch_size:]) self.batch_total_loss = (cross_ent_loss, adv_cross_ent_loss) loss = (weighted_kl_loss + cross_ent_loss / num_tokens + self.max_margin_weight * tf.reduce_mean(max_margin_loss) ) # max_margin_loss shape=[self.half_batch_size] else: cross_ent_loss = tf.reduce_sum(total_loss) self.batch_total_loss = cross_ent_loss loss = weighted_kl_loss + cross_ent_loss / num_tokens self.batch_num_tokens = num_tokens grads = compute_grads(loss, optimizer, self.trainable_variables) # First context input will be repeated in the next batch, so we ignore it. assign_context_input_op = tf.assign( context_input_acc, context_input[:, 1:, :], validate_shape=False) # shape will be differnt on axis=1 with tf.control_dependencies([ assign_context_input_op ]): # make sure we update context_input_acc self.apply_gradients_op = apply_grads( optimizer, [grads], clipping_threshold=self.clipping_threshold, global_step=self.global_step) # Just for control dependencies self.batch_sample_ids_beam = tf.identity(sample_ids) self.batch_final_lengths_beam = tf.identity(final_lengths)
def __init__(self, batch_size, shared_vocab_size_politeness, new_vocab_size_politeness, shared_vocab_size_movie, new_vocab_size_movie, embedding_word2vec_politeness, embedding_word2vec_movie, embedding_size, hidden_size, num_layers, max_iterations, start_token, end_token, unk_indices, attention_size=512, attention_layer_size=256, beam_width=10, length_penalty_weight=1.0, gpu_start_index=0, num_gpus=1, # set to 1 when testing learning_rate=0.0001, clipping_threshold=5.0, # str_input=False, vocab=None, feed_final_state=True, sampling_prob=0.25, monotonic_attention=False, proj_embedding=True, continuous_label=False, output_layer=False, clip_loss=False, clip_ratio=0.2, decay_ratio=0.999): """ Args: num_layers: set it to 2. max_iterations: should be set to max_iterations """ self.batch_size = batch_size self.batch_size_per_gpu = batch_size // num_gpus self.shared_vocab_size_politeness = shared_vocab_size_politeness self.new_vocab_size_politeness = new_vocab_size_politeness self.shared_vocab_size_movie = shared_vocab_size_movie self.new_vocab_size_movie = new_vocab_size_movie self.vocab_size = (1 + # +1 for UNK_TOKEN shared_vocab_size_politeness + new_vocab_size_politeness + shared_vocab_size_movie + new_vocab_size_movie) self.embedding_word2vec_politeness = embedding_word2vec_politeness self.embedding_word2vec_movie = embedding_word2vec_movie self.embedding_size = embedding_size self.hidden_size = hidden_size assert self.hidden_size % 2 == 0 self.num_layers = num_layers self.max_iterations = max_iterations self.start_token = start_token self.end_token = end_token self.unk_indices = unk_indices self.attention_size = attention_size self.attention_layer_size = attention_layer_size self.beam_width = beam_width self.length_penalty_weight = length_penalty_weight self.clipping_threshold = clipping_threshold # self.str_input = str_input # self.vocab = vocab self.feed_final_state = feed_final_state self.sampling_prob = sampling_prob # probability for scheduled sampling self.monotonic_attention = monotonic_attention self.proj_embedding = proj_embedding self.continuous_label = continuous_label self.output_layer = True if output_layer else None self.clip_loss = clip_loss self.clip_ratio = clip_ratio self.decay_ratio = decay_ratio self.trainable_variables = [] with tf.variable_scope("seq2seq", reuse=tf.AUTO_REUSE): with tf.device('/cpu:0'): self.create_placeholders() self.avg_loss = tf.get_variable( "avg_loss", shape=[], initializer=tf.constant_initializer(4.0, dtype=tf.float32), trainable=False) # if self.str_input: # source = self.indexed_source # target = self.indexed_target # else: source = self.source target = self.target # Note: Make sure batch size can be evenly divided by num_gpus split = functools.partial(tf.split, num_or_size_splits=num_gpus, axis=0) [source_lst, source_length_lst, target_lst, target_length_lst, start_tokens_lst] = [ split(tensor) for tensor in [source, self.source_length, target, self.target_length, self.start_tokens]] optimizer = tf.train.AdamOptimizer(learning_rate) sample_ids_beam_lst = [] final_lengths_beam_lst = [] num_tokens_lst = [] total_losses = [] tower_grads = [] for i in xrange(num_gpus): source_max_length = tf.reduce_max(source_length_lst[i]) target_max_length = tf.reduce_max(target_length_lst[i]) (num_tokens, total_loss, grads, sample_ids_beam, final_lengths_beam) = self.one_iteration( source_lst[i][:, :source_max_length], source_length_lst[i], target_lst[i][:, :target_max_length], target_length_lst[i], start_tokens_lst[i], optimizer, gpu_index=(gpu_start_index + i)) sample_ids_beam_lst.append( pad_and_truncate(sample_ids_beam, final_lengths_beam, self.max_iterations)) final_lengths_beam_lst.append(final_lengths_beam) num_tokens_lst.append(num_tokens) total_losses.append(total_loss) tower_grads.append(grads) with tf.device('/cpu:0'): self.batch_total_loss = tf.add_n(total_losses) self.batch_num_tokens = tf.add_n(num_tokens_lst) # Concat sample ids and their respective lengths self.batch_sample_ids_beam = tf.concat( sample_ids_beam_lst, axis=0) self.batch_final_lengths_beam = tf.concat( final_lengths_beam_lst, axis=0) self.apply_gradients_op = apply_grads( optimizer, tower_grads)
def __init__(self, batch_size, vocab_size, embedding_size, hidden_size_encoder, hidden_size_context, hidden_size_decoder, dim_z, max_iterations, max_dialogue_length, start_token, end_token, unk_indices, num_layers_encoder=1, num_layers_context=1, num_layers_decoder=1, attention_size=512, attention_layer_size=256, beam_width=10, length_penalty_weight=1.0, gpu_start_index=0, num_gpus=1, # set to 1 when testing learning_rate=0.001, clipping_threshold=5.0, truncated=True): self.batch_size = batch_size self.batch_size_per_gpu = batch_size // num_gpus self.vocab_size = vocab_size self.embedding_size = embedding_size self.hidden_size_encoder = hidden_size_encoder self.hidden_size_context = hidden_size_context self.hidden_size_decoder = hidden_size_decoder assert self.hidden_size_context * 2 == self.hidden_size_decoder self.dim_z = dim_z self.dense = functools.partial(tf.layers.dense, units=self.dim_z, use_bias=True) self.max_iterations = max_iterations self.max_dialogue_length = max_dialogue_length assert self.max_dialogue_length > 0 self.start_tokens = [start_token] * self.batch_size_per_gpu self.end_token = end_token self.unk_indices = unk_indices self.num_layers_encoder = num_layers_encoder self.num_layers_context = num_layers_context self.num_layers_decoder = num_layers_decoder self.attention_size = attention_size self.attention_layer_size = attention_layer_size self.beam_width = beam_width self.length_penalty_weight = length_penalty_weight self.num_gpus = num_gpus self.clipping_threshold = clipping_threshold self.truncated = truncated """ context_input_acc put context inputs in reverse order. • When adding inputs, just concat from left • When using, we do not need to reverse it back because we are using bidirectional-lstm Note: If we don't use bidirectional-lstm, then we will need tf.reverse_sequence() """ self.trainable_variables = [] with tf.variable_scope("seq2seq"): with tf.device('/cpu:0'): self.create_placeholders() context_input_acc = tf.get_variable( "context_input_acc", initializer=tf.zeros([self.batch_size, 2, self.hidden_size_encoder * 2], dtype=tf.float32), trainable=False) # max_num_turns better not be less than 2, since we may just lose a whole dimension (i.e., axis=1)? max_num_turns = tf.maximum(tf.reduce_max(self.start_turn_index), 2) context_input_mask = tf.tile( tf.reshape(tf.greater(self.start_turn_index, 0), [self.batch_size, 1, 1]), # expand two dims [1, max_num_turns, self.hidden_size_encoder * 2]) # This multiplication resets context input that have start_turn_index == 0 previous_context_input = context_input_acc[:, :max_num_turns, :] * tf.cast(context_input_mask, tf.float32) # Note: Make sure batch_size can be evenly divided by num_gpus [dialogue_lst, turn_length_lst, previous_context_input_lst, start_turn_index_lst] = [ tf.split(tensor, self.num_gpus, axis=0) for tensor in [self.dialogue, self.turn_length, previous_context_input, #cannot be less than 2, otherwise tf.map_fn will give error. self.start_turn_index]] optimizer = tf.train.AdamOptimizer(learning_rate) context_input_lst = [] sample_ids_beam_lst = [] final_lengths_beam_lst = [] num_tokens_lst = [] total_losses = [] tower_grads = [] for i in xrange(num_gpus): with tf.device("/gpu:%d" % (gpu_start_index + i)): (total_loss, num_tokens, context_input, sample_ids_beam, final_lengths_beam) = self.one_iteration( dialogue_lst[i], turn_length_lst[i], previous_context_input_lst[i], start_turn_index_lst[i], optimizer) # first turn will be repeated in the next batch, so we skip it context_input_lst.append(context_input[:, 1:, :]) sample_ids_beam_lst.append(sample_ids_beam) final_lengths_beam_lst.append(final_lengths_beam) grads = compute_grads( total_loss / num_tokens, optimizer, self.trainable_variables) tower_grads.append(grads) total_losses.append(total_loss) num_tokens_lst.append(num_tokens) with tf.device('/cpu:0'): context_input_concat = tf.concat(context_input_lst, axis=0) assign_context_input_op = tf.assign( context_input_acc, context_input_concat, validate_shape=False) # shape will be differnt on axis=1 with tf.control_dependencies([assign_context_input_op]): # make sure we update context_input_acc # Concat sample ids and their respective lengths self.batch_sample_ids_beam = tf.concat( sample_ids_beam_lst, axis=0) self.batch_final_lengths_beam = tf.concat( final_lengths_beam_lst, axis=0) self.batch_total_loss = tf.add_n(total_losses) self.batch_num_tokens = tf.add_n(num_tokens_lst) self.apply_gradients_op = apply_grads( optimizer, tower_grads, clipping_threshold=self.clipping_threshold)