def train(num_training_iterations, report_interval, reduce_learning_rate_interval): """Run the training of the deep LSTM model on tiny shakespeare.""" dataset_train = dataset_shakespeare.TinyShakespeareDataset( num_steps=FLAGS.truncation_length, batch_size=FLAGS.batch_size, subset="train", random=True, name="shake_train") dataset_valid = dataset_shakespeare.TinyShakespeareDataset( num_steps=FLAGS.truncation_length, batch_size=FLAGS.batch_size, subset="valid", random=False, name="shake_valid") dataset_test = dataset_shakespeare.TinyShakespeareDataset( num_steps=FLAGS.truncation_length, batch_size=FLAGS.batch_size, subset="test", random=False, name="shake_test") model = TextModel( num_embedding=FLAGS.num_embedding, num_hidden=FLAGS.num_hidden, lstm_depth=FLAGS.lstm_depth, output_size=dataset_valid.vocab_size, use_dynamic_rnn=True, use_skip_connections=True) # Build the training model and get the training loss. train_input_sequence, train_target_sequence = dataset_train() train_output_sequence_logits, train_final_state = model(train_input_sequence) # pylint: disable=not-callable train_loss = dataset_train.cost(train_output_sequence_logits, train_target_sequence) # Get the validation loss. valid_input_sequence, valid_target_sequence = dataset_valid() valid_output_sequence_logits, _ = model(valid_input_sequence) # pylint: disable=not-callable valid_loss = dataset_valid.cost(valid_output_sequence_logits, valid_target_sequence) # Get the test loss. test_input_sequence, test_target_sequence = dataset_test() test_output_sequence_logits, _ = model(test_input_sequence) # pylint: disable=not-callable test_loss = dataset_test.cost(test_output_sequence_logits, test_target_sequence) # Build graph to sample some strings during training. initial_logits = train_output_sequence_logits[FLAGS.truncation_length - 1] train_generated_string = model.generate_string( initial_logits=initial_logits, initial_state=train_final_state, sequence_length=FLAGS.sample_length) # Set up optimizer with global norm clipping. trainable_variables = tf.trainable_variables() grads, _ = tf.clip_by_global_norm( tf.gradients(train_loss, trainable_variables), FLAGS.max_grad_norm) learning_rate = tf.get_variable( "learning_rate", shape=[], dtype=tf.float32, initializer=tf.constant_initializer(FLAGS.learning_rate), trainable=False) reduce_learning_rate = learning_rate.assign( learning_rate * FLAGS.reduce_learning_rate_multiplier) global_step = tf.get_variable( name="global_step", shape=[], dtype=tf.int64, initializer=tf.zeros_initializer(), trainable=False, collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP]) optimizer = tf.train.AdamOptimizer( learning_rate, epsilon=FLAGS.optimizer_epsilon) train_step = optimizer.apply_gradients( zip(grads, trainable_variables), global_step=global_step) saver = tf.train.Saver() hooks = [ tf.train.CheckpointSaverHook( checkpoint_dir=FLAGS.checkpoint_dir, save_steps=FLAGS.checkpoint_interval, saver=saver) ] # Train. with tf.train.SingularMonitoredSession( hooks=hooks, checkpoint_dir=FLAGS.checkpoint_dir) as sess: start_iteration = sess.run(global_step) for train_iteration in xrange(start_iteration, num_training_iterations): if (train_iteration + 1) % report_interval == 0: train_loss_v, valid_loss_v, _ = sess.run( (train_loss, valid_loss, train_step)) train_generated_string_v = sess.run(train_generated_string) train_generated_string_human = dataset_train.to_human_readable( (train_generated_string_v, 0), indices=[0]) tf.logging.info("%d: Training loss %f. Validation loss %f. Sample = %s", train_iteration, train_loss_v, valid_loss_v, train_generated_string_human) else: train_loss_v, _ = sess.run((train_loss, train_step)) tf.logging.info("%d: Training loss %f.", train_iteration, train_loss_v) if (train_iteration + 1) % reduce_learning_rate_interval == 0: sess.run(reduce_learning_rate) tf.logging.info("Reducing learning rate.") test_loss = sess.run(test_loss) tf.logging.info("Test loss %f", test_loss)
def build_graph(lstm_depth=3, batch_size=32, num_embedding=32, num_hidden=128, truncation_length=64, sample_length=1000, max_grad_norm=5, initial_learning_rate=0.1, reduce_learning_rate_multiplier=0.1, optimizer_epsilon=0.01): """Constructs the computation graph.""" # Get datasets. dataset_train = dataset_shakespeare.TinyShakespeareDataset( num_steps=truncation_length, batch_size=batch_size, subset="train", random=True, name="shake_train") dataset_valid = dataset_shakespeare.TinyShakespeareDataset( num_steps=truncation_length, batch_size=batch_size, subset="valid", random=False, name="shake_valid") dataset_test = dataset_shakespeare.TinyShakespeareDataset( num_steps=truncation_length, batch_size=batch_size, subset="test", random=False, name="shake_test") # Define model. model = TextModel( num_embedding=num_embedding, num_hidden=num_hidden, lstm_depth=lstm_depth, output_size=dataset_valid.vocab_size, use_dynamic_rnn=True, use_skip_connections=True) # Get the training loss. train_input_sequence, train_target_sequence = dataset_train() train_output_sequence_logits, train_final_state = model(train_input_sequence) # pylint: disable=not-callable train_loss = dataset_train.cost(train_output_sequence_logits, train_target_sequence) # Get the validation loss. valid_input_sequence, valid_target_sequence = dataset_valid() valid_output_sequence_logits, _ = model(valid_input_sequence) # pylint: disable=not-callable valid_loss = dataset_valid.cost(valid_output_sequence_logits, valid_target_sequence) # Get the test loss. test_input_sequence, test_target_sequence = dataset_test() test_output_sequence_logits, _ = model(test_input_sequence) # pylint: disable=not-callable test_loss = dataset_test.cost(test_output_sequence_logits, test_target_sequence) # Build graph to sample some strings during training. initial_logits = train_output_sequence_logits[truncation_length - 1] train_generated_string = model.generate_string( initial_logits=initial_logits, initial_state=train_final_state, sequence_length=sample_length) # Set up global norm clipping of gradients. trainable_variables = tf.trainable_variables() grads, _ = tf.clip_by_global_norm( tf.gradients(train_loss, trainable_variables), max_grad_norm) # Get learning rate and define annealing. learning_rate = tf.get_variable( "learning_rate", shape=[], dtype=tf.float32, initializer=tf.constant_initializer(initial_learning_rate), trainable=False) reduce_learning_rate = learning_rate.assign( learning_rate * reduce_learning_rate_multiplier) # Get training step counter. global_step = tf.get_variable( name="global_step", shape=[], dtype=tf.int64, initializer=tf.zeros_initializer(), trainable=False, collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP]) # Define optimizer and training step. optimizer = tf.train.AdamOptimizer( learning_rate, epsilon=optimizer_epsilon) train_step = optimizer.apply_gradients( zip(grads, trainable_variables), global_step=global_step) graph_tensors = { "train_loss": train_loss, "valid_loss": valid_loss, "test_loss": test_loss, "train_generated_string": train_generated_string, "reduce_learning_rate": reduce_learning_rate, "global_step": global_step, "train_step": train_step } # Return dataset_train for translation to human readable text. return graph_tensors, dataset_train