from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

SEED = 67

import time
import random; random.seed(SEED)
import numpy as np; np.random.seed(SEED)
import tensorflow as tf; tf.set_random_seed(SEED)

from entity_networks.model import Model
from entity_networks.dataset import Dataset

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_integer('batch_size', 32, 'Batch size.')
tf.app.flags.DEFINE_integer('num_epochs', 200, 'Number of training epochs.')
tf.app.flags.DEFINE_string('logdir', 'logs/{}'.format(int(time.time())), 'Log directory.')

def main(_):
    with tf.device('/cpu:0'):
        dataset_train = Dataset(
            filenames=['datasets/processed/qa1_single-supporting-fact_train.tfrecords'],
            batch_size=FLAGS.batch_size,
            shuffle=True)
        dataset_test = Dataset(
            filenames=['datasets/processed/qa1_single-supporting-fact_test.tfrecords'],
            batch_size=FLAGS.batch_size,
            shuffle=False)

    with tf.variable_scope('model'):
        model_train = Model(dataset_train, is_training=True)

    with tf.variable_scope('model', reuse=True):
        model_test = Model(dataset_test, is_training=False)

    print('Training model with {} parameters'.format(model_train.num_parameters))

    max_steps = FLAGS.num_epochs*dataset_train.size

    tf.contrib.learn.train(
        graph=tf.get_default_graph(),
        output_dir=FLAGS.logdir,
        train_op=model_train.train_op,
        loss_op=model_train.loss,
        supervisor_save_summaries_steps=1,
        log_every_steps=1,
        max_steps=max_steps)

if __name__ == '__main__':
    tf.app.run()

@@ -334,16 +334,69 @@
  Dataset
+%0Afrom entity_networks.monitors import ProgressMonitor
 %0A%0AFLAGS 
@@ -1092,25 +1092,25 @@
 able_scope('
-m
+M
 odel'):%0A    
@@ -1163,24 +1163,59 @@
 ning=True)%0A%0A
+    # TODO: use validation monitor%0A
     with tf.
@@ -1230,17 +1230,17 @@
 _scope('
-m
+M
 odel', r
@@ -1308,24 +1308,76 @@
 ing=False)%0A%0A
+    print('Saving logs to %7B%7D'.format(FLAGS.logdir))%0A
     print('T
@@ -1494,20 +1494,485 @@
 t_train.
-size
+num_batches%0A    monitors = %5B%0A        tf.contrib.learn.monitors.SummarySaver(%0A            summary_op=tf.merge_all_summaries(),%0A            save_steps=1,%0A            output_dir=FLAGS.logdir),%0A        ProgressMonitor(tensor_names=%7B%0A            'Loss (Train)': model_train.loss.name,%0A            'Loss (Test)': model_test.loss.name,%0A            'LR': model_train.learning_rate.name,%0A        %7D, every_n_steps=1000, first_n_steps=0), # TODO: fix every N: logs too often%0A    %5D
 %0A%0A    tf
@@ -2137,16 +2137,16 @@
 n.loss,%0A
-
         
@@ -2172,32 +2172,59 @@
 maries_steps=1,%0A
+        monitors=monitors,%0A
         log_ever

Add monitors