Make tensorboard optional

Author: tuzhucheng

.gitignore

@@ -1,4 +1,5 @@
 __pycache__/
 saved_models/
+.idea/
 *.swp
 *.pyc

README.md

@@ -48,4 +48,4 @@ python main.py --help
 
 ## Dependencies
 
-The model is written in PyTorch. We make use of https://github.com/lanpa/tensorboard-pytorch to connect to [TensorBoard](https://github.com/tensorflow/tensorboard) to visualize the training process.
+The model is written in PyTorch. We optionally make use of https://github.com/lanpa/tensorboard-pytorch to connect to [TensorBoard](https://github.com/tensorflow/tensorboard) to visualize the training process. Just add `--tensorboard` to enable.

main.py

@@ -55,6 +55,7 @@
     parser.add_argument('--padding', action='store_true', default=False, help='use padding for convolution (default: False)')
     parser.add_argument('--dropout', type=float, default=0.5, help='dropout probability (default: 0.5)')
     parser.add_argument('--seed', type=int, default=1, help='random seed (default: 1)')
+    parser.add_argument('--tensorboard', action='store_true', default=False, help='use TensorBoard to visualize training (default: false)')
     parser.add_argument('--run-label', type=str, help='label to describe run')
     args = parser.parse_args()
     args.cuda = not args.no_cuda and torch.cuda.is_available()

train.py

@@ -1,7 +1,6 @@
 import math
 import time
 
-from tensorboardX import SummaryWriter
 import torch
 import torch.nn.functional as F
 from torch.autograd import Variable
@@ -50,11 +49,14 @@ def __init__(self, model, train_loader, trainer_config, train_evaluator, test_ev
         self.model_outfile = trainer_config['model_outfile']
         self.lr_reduce_factor = trainer_config['lr_reduce_factor']
         self.patience = trainer_config['patience']
+        self.use_tensorboard = trainer_config['tensorboard']
+        if self.use_tensorboard:
+            from tensorboardX import SummaryWriter
+            self.writer = SummaryWriter(log_dir=None, comment='' if trainer_config['run_label'] is None else trainer_config['run_label'])
 
         self.train_evaluator = train_evaluator
         self.test_evaluator = test_evaluator
         self.dev_evaluator = dev_evaluator
-        self.writer = SummaryWriter(log_dir=None, comment=trainer_config['run_label'])
 
     def evaluate(self, evaluator, dataset_name):
         scores, metric_names = evaluator.get_scores()
@@ -94,7 +96,10 @@ def train_epoch(self, epoch):
                     len(self.train_loader.dataset) if not self.sample else self.sample,
                     100. * batch_idx / (len(self.train_loader) if not self.sample else math.ceil(self.sample / self.batch_size)), loss.data[0])
                 )
-        self.writer.add_scalar('sick/train/kl_div_loss', total_loss, epoch)
+
+        if self.use_tensorboard:
+            self.writer.add_scalar('sick/train/kl_div_loss', total_loss, epoch)
+
         return total_loss
 
     def train(self, epochs):
@@ -109,9 +114,12 @@ def train(self, epochs):
 
             dev_scores = self.evaluate(self.dev_evaluator, 'dev')
             new_loss = dev_scores[2]
-            self.writer.add_scalar('sick/lr', self.optimizer.param_groups[0]['lr'], epoch)
-            self.writer.add_scalar('sick/dev/pearson_r', dev_scores[0], epoch)
-            self.writer.add_scalar('sick/dev/kl_div_loss', new_loss, epoch)
+
+            if self.use_tensorboard:
+                self.writer.add_scalar('sick/lr', self.optimizer.param_groups[0]['lr'], epoch)
+                self.writer.add_scalar('sick/dev/pearson_r', dev_scores[0], epoch)
+                self.writer.add_scalar('sick/dev/kl_div_loss', new_loss, epoch)
+
             end = time.time()
             duration = end - start
             logger.info('Epoch {} finished in {:.2f} minutes'.format(epoch, duration / 60))
@@ -170,7 +178,10 @@ def train_epoch(self, epoch):
                 )
 
         self.evaluate(self.train_evaluator, 'train')
-        self.writer.add_scalar('msrvid/train/kl_div_loss', total_loss, epoch)
+
+        if self.use_tensorboard:
+            self.writer.add_scalar('msrvid/train/kl_div_loss', total_loss, epoch)
+
         return left_out_val_a, left_out_val_b, left_out_ext_feats, left_out_val_labels
 
     def train(self, epochs):
@@ -196,12 +207,17 @@ def train(self, epochs):
             predictions = predictions.cpu().numpy()
             true_labels = true_labels.cpu().numpy()
             pearson_r = pearsonr(predictions, true_labels)[0]
-            self.writer.add_scalar('msrvid/dev/pearson_r', pearson_r, epoch)
+
+            if self.use_tensorboard:
+                self.writer.add_scalar('msrvid/dev/pearson_r', pearson_r, epoch)
+
             for param_group in self.optimizer.param_groups:
                 logger.info('Validation size: %s Pearson\'s r: %s', output.size()[0], pearson_r)
                 logger.info('Learning rate: %s', param_group['lr'])
-                self.writer.add_scalar('msrvid/lr', param_group['lr'], epoch)
-                self.writer.add_scalar('msrvid/dev/kl_div_loss', val_kl_div_loss, epoch)
+
+                if self.use_tensorboard:
+                    self.writer.add_scalar('msrvid/lr', param_group['lr'], epoch)
+                    self.writer.add_scalar('msrvid/dev/kl_div_loss', val_kl_div_loss, epoch)
                 break
             scheduler.step(pearson_r)