add tqdm and tensorboard support on example (#35)

.gitignore

@@ -14,3 +14,5 @@ a.out
 .ipynb_checkpoints
 benchmark_xyz
 *.pyc
+*checkpoint*
+/runs
\ No newline at end of file

Dockerfile

 FROM zasdfgbnm/pytorch-master
-RUN pacman -Sy --noconfirm python-sphinx python2-sphinx flake8
+RUN pacman -Sy --noconfirm python-sphinx python2-sphinx python-tqdm python2-tqdm flake8
+RUN pip install tensorboardX && pip2 install tensorboardX
 COPY . /torchani
 RUN cd torchani && pip install .
 RUN cd torchani && pip2 install .

examples/nnp_training.py

@@ -3,12 +3,20 @@ import torch
 import ignite
 import torchani
 import model
+import tqdm
+import timeit
+import tensorboardX
+import math
 
 chunk_size = 256
 batch_chunks = 4
 dataset_path = sys.argv[1]
 dataset_checkpoint = 'dataset-checkpoint.dat'
 model_checkpoint = 'checkpoint.pt'
+max_epochs = 10
+
+writer = tensorboardX.SummaryWriter()
+start = timeit.default_timer()
 
 shift_energy = torchani.EnergyShifter()
 training, validation, testing = torchani.data.load_or_create(
@@ -40,26 +48,44 @@ evaluator = ignite.engine.create_supervised_evaluator(container, metrics={
     })
 
 
+def hartree2kcal(x):
+    return 627.509 * x
+
+
+@trainer.on(ignite.engine.Events.EPOCH_STARTED)
+def init_tqdm(trainer):
+    trainer.state.tqdm = tqdm.tqdm(total=len(training), desc='epoch')
+
+
 @trainer.on(ignite.engine.Events.ITERATION_COMPLETED)
-def log_training_loss(trainer):
-    print("Epoch[{}] Loss: {:.2f}".format(trainer.state.epoch,
-                                          trainer.state.output))
+def update_tqdm(trainer):
+    trainer.state.tqdm.update(1)
 
 
 @trainer.on(ignite.engine.Events.EPOCH_COMPLETED)
-def log_training_results(trainer):
-    evaluator.run(training)
-    metrics = evaluator.state.metrics
-    print("Training Results - Epoch: {}  RMSE: {:.2f}"
-          .format(trainer.state.epoch, metrics['RMSE']))
+def finalize_tqdm(trainer):
+    trainer.state.tqdm.close()
 
 
-@trainer.on(ignite.engine.Events.EPOCH_COMPLETED)
+@trainer.on(ignite.engine.Events.EPOCH_STARTED)
 def log_validation_results(trainer):
     evaluator.run(validation)
     metrics = evaluator.state.metrics
-    print("Validation Results - Epoch: {}  RMSE: {:.2f}"
-          .format(trainer.state.epoch, metrics['RMSE']))
+    rmse = hartree2kcal(metrics['RMSE'])
+    writer.add_scalar('validation_rmse_vs_epoch', rmse, trainer.state.epoch)
+
+
+@trainer.on(ignite.engine.Events.EPOCH_STARTED)
+def log_time(trainer):
+    elapsed = round(timeit.default_timer() - start, 2)
+    writer.add_scalar('time_vs_epoch', elapsed, trainer.state.epoch)
+
+
+@trainer.on(ignite.engine.Events.ITERATION_COMPLETED)
+def log_loss_and_time(trainer):
+    iteration = trainer.state.iteration
+    rmse = hartree2kcal(math.sqrt(trainer.state.output))
+    writer.add_scalar('training_rmse_vs_iteration', rmse, iteration)
 
 
-trainer.run(training, max_epochs=10)
+trainer.run(training, max_epochs=max_epochs)

examples/training-benchmark.py

@@ -4,6 +4,7 @@ import ignite
 import torchani
 import timeit
 import model
+import tqdm
 
 chunk_size = 256
 batch_chunks = 4
@@ -33,6 +34,22 @@ optimizer = torch.optim.Adam(nnp.parameters())
 trainer = ignite.engine.create_supervised_trainer(
     container, optimizer, torchani.ignite.energy_mse_loss)
 
+
+@trainer.on(ignite.engine.Events.EPOCH_STARTED)
+def init_tqdm(trainer):
+    trainer.state.tqdm = tqdm.tqdm(total=len(dataloader), desc='epoch')
+
+
+@trainer.on(ignite.engine.Events.ITERATION_COMPLETED)
+def update_tqdm(trainer):
+    trainer.state.tqdm.update(1)
+
+
+@trainer.on(ignite.engine.Events.EPOCH_COMPLETED)
+def finalize_tqdm(trainer):
+    trainer.state.tqdm.close()
+
+
 start = timeit.default_timer()
 trainer.run(dataloader, max_epochs=1)
 elapsed = round(timeit.default_timer() - start, 2)