From 857cd43934c651e9b4c2c4819692284ad4e944be Mon Sep 17 00:00:00 2001 From: Michael Tu Date: Thu, 5 Oct 2017 19:15:51 -0400 Subject: [PATCH] MP-CNN: Optionally Visualize Training for Debugging (#67) * MP-CNN: optionally support TensorBoard for learning curve visualization * MP-CNN: bug with SummaryWriter comment * MP-CNN: add instructions on how to setup tensorboard --- mp_cnn/README.md | 9 +++++++ mp_cnn/main.py | 15 +++++++++++- mp_cnn/train.py | 62 +++++++++++++++++++++++++++++++++--------------- 3 files changed, 66 insertions(+), 20 deletions(-) diff --git a/mp_cnn/README.md b/mp_cnn/README.md index 1e0b78b..3875378 100644 --- a/mp_cnn/README.md +++ b/mp_cnn/README.md @@ -47,3 +47,12 @@ To see all options available, use ``` python main.py --help ``` + +## Optional Dependencies + +To optionally visualize the learning curve during training, we make use of https://github.com/lanpa/tensorboard-pytorch to connect to [TensorBoard](https://github.com/tensorflow/tensorboard). These projects require TensorFlow as a dependency, so you need to install TensorFlow before running the commands below. After these are installed, just add `--tensorboard` when running `main.py` and open TensorBoard in the browser. + +```sh +pip install tensorboardX +pip install tensorflow-tensorboard +``` diff --git a/mp_cnn/main.py b/mp_cnn/main.py index 9312dfc..d265b27 100644 --- a/mp_cnn/main.py +++ b/mp_cnn/main.py @@ -47,6 +47,8 @@ if __name__ == '__main__': parser.add_argument('--hidden-units', type=int, default=150, help='number of hidden units in each of the two hidden layers (default: 150)') 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() @@ -73,7 +75,18 @@ if __name__ == '__main__': test_evaluator = MPCNNEvaluatorFactory.get_evaluator(args.dataset, model, test_loader, args.batch_size, args.cuda) dev_evaluator = MPCNNEvaluatorFactory.get_evaluator(args.dataset, model, dev_loader, args.batch_size, args.cuda) - trainer = MPCNNTrainerFactory.get_trainer(args.dataset, model, optimizer, train_loader, args.batch_size, args.sample, args.log_interval, args.model_outfile, args.lr_reduce_factor, args.patience, train_evaluator, test_evaluator, dev_evaluator) + trainer_config = { + 'optimizer': optimizer, + 'batch_size': args.batch_size, + 'sample': args.sample, + 'log_interval': args.log_interval, + 'model_outfile': args.model_outfile, + 'lr_reduce_factor': args.lr_reduce_factor, + 'patience': args.patience, + 'tensorboard': args.tensorboard, + 'run_label': args.run_label + } + trainer = MPCNNTrainerFactory.get_trainer(args.dataset, model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator) if not args.skip_training: total_params = 0 diff --git a/mp_cnn/train.py b/mp_cnn/train.py index b8a4322..62e93cb 100644 --- a/mp_cnn/train.py +++ b/mp_cnn/train.py @@ -24,11 +24,11 @@ class MPCNNTrainerFactory(object): Get the corresponding Trainer class for a particular dataset. """ @staticmethod - def get_trainer(dataset_name, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator=None): + def get_trainer(dataset_name, model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator=None): if dataset_name == 'sick': - return SICKTrainer(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator) + return SICKTrainer(model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator) elif dataset_name == 'msrvid': - return MSRVIDTrainer(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile,lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator) + return MSRVIDTrainer(model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator) else: raise ValueError('{} is not a valid dataset.'.format(dataset_name)) @@ -39,16 +39,21 @@ class Trainer(object): Abstraction for training a model on a Dataset. """ - def __init__(self, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator=None): + def __init__(self, model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator=None): self.model = model - self.optimizer = optimizer + self.optimizer = trainer_config['optimizer'] self.train_loader = train_loader - self.batch_size = batch_size - self.sample = sample - self.log_interval = log_interval - self.model_outfile = model_outfile - self.lr_reduce_factor = lr_reduce_factor - self.patience = patience + self.batch_size = trainer_config['batch_size'] + self.sample = trainer_config['sample'] + self.log_interval = trainer_config['log_interval'] + 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 @@ -69,8 +74,8 @@ class Trainer(object): class SICKTrainer(Trainer): - def __init__(self, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator=None): - super(SICKTrainer, self).__init__(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator) + def __init__(self, model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator=None): + super(SICKTrainer, self).__init__(model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator) def train_epoch(self, epoch): self.model.train() @@ -92,7 +97,8 @@ class SICKTrainer(Trainer): 100. * batch_idx / (len(self.train_loader) if not self.sample else math.ceil(self.sample / self.batch_size)), loss.data[0]) ) - del loss, output + if self.use_tensorboard: + self.writer.add_scalar('sick/train/kl_div_loss', total_loss, epoch) return total_loss @@ -108,6 +114,12 @@ class SICKTrainer(Trainer): dev_scores = self.evaluate(self.dev_evaluator, 'dev') new_loss = dev_scores[2] + + 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)) @@ -129,8 +141,8 @@ class SICKTrainer(Trainer): class MSRVIDTrainer(Trainer): - def __init__(self, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator=None): - super(MSRVIDTrainer, self).__init__(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, lr_reduce_factor, patience, train_evaluator, test_evaluator, dev_evaluator) + def __init__(self, model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator=None): + super(MSRVIDTrainer, self).__init__(model, train_loader, trainer_config, train_evaluator, test_evaluator, dev_evaluator) def train_epoch(self, epoch): self.model.train() @@ -141,7 +153,7 @@ class MSRVIDTrainer(Trainer): left_out_val_a, left_out_val_b = [], [] left_out_ext_feats = [] left_out_val_labels = [] - + total_loss = 0 for batch_idx, (sentences, labels) in enumerate(self.train_loader): sent_a, sent_b = Variable(sentences['a']), Variable(sentences['b']) ext_feats = Variable(sentences['ext_feats']) @@ -154,6 +166,7 @@ class MSRVIDTrainer(Trainer): self.optimizer.zero_grad() output = self.model(sent_a, sent_b, ext_feats) loss = F.kl_div(output, labels) + total_loss += loss.data[0] loss.backward() self.optimizer.step() if batch_idx % self.log_interval == 0: @@ -163,9 +176,11 @@ class MSRVIDTrainer(Trainer): 100. * batch_idx / (len(self.train_loader) if not self.sample else math.ceil(self.sample / self.batch_size)), loss.data[0]) ) - del loss, output - self.evaluate(self.train_evaluator, 'train') + + 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): @@ -191,10 +206,19 @@ class MSRVIDTrainer(Trainer): predictions = predictions.cpu().numpy() true_labels = true_labels.cpu().numpy() pearson_r = pearsonr(predictions, true_labels)[0] + + 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']) + + 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) end = time.time()