added gradient clipping

docs/Examples.md

@@ -1,23 +1,32 @@
+#### Template model definition
+In 99% of cases you want to just copy this template to start a new lightningModule and change the core of what your model is actually trying to do.
+
+``` {.python}
+import os
+from collections import OrderedDict
 import torch.nn as nn
-import numpy as np
-from pytorch_lightning.root_module.root_module import LightningModule
-from test_tube import HyperOptArgumentParser
 from torchvision.datasets import MNIST
 import torchvision.transforms as transforms
 import torch
 import torch.nn.functional as F
-import os, pdb
+from test_tube import HyperOptArgumentParser
-from collections import OrderedDict
+from torch import optim
+
+from pytorch_lightning.root_module.root_module import LightningModule
 
 
-class ExampleModel(LightningModule):
+class LightningTemplateModel(LightningModule):
     """
     Sample model to show how to define a template
     """
 
     def __init__(self, hparams):
+        """
+        Pass in parsed HyperOptArgumentParser to the model
+        :param hparams:
+        """
         # init superclass
-        super(ExampleModel, self).__init__(hparams)
+        super(LightningTemplateModel, self).__init__(hparams)
 
         self.batch_size = hparams.batch_size
 
@@ -42,6 +51,11 @@ class ExampleModel(LightningModule):
     # TRAINING
     # ---------------------
     def forward(self, x):
+        """
+        No special modification required for lightning, define as you normally would
+        :param x:
+        :return:
+        """
 
         x = self.c_d1(x)
         x = torch.tanh(x)
@@ -59,7 +73,7 @@ class ExampleModel(LightningModule):
 
     def training_step(self, data_batch, batch_i):
         """
-        Called inside the training loop
+        Lightning calls this inside the training loop
         :param data_batch:
         :return:
         """
@@ -79,7 +93,7 @@ class ExampleModel(LightningModule):
 
     def validation_step(self, data_batch, batch_i):
         """
-        Called inside the validation loop
+        Lightning calls this inside the validation loop
         :param data_batch:
         :return:
         """
@@ -99,7 +113,6 @@ class ExampleModel(LightningModule):
         })
         return output
 
-
     def validation_end(self, outputs):
         """
         Called at the end of validation to aggregate outputs
@@ -139,9 +152,8 @@ class ExampleModel(LightningModule):
         return whatever optimizers we want here
         :return: list of optimizers
         """
-        optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
-        self.optimizers = [optimizer]
+        return [optimizer]
-        return self.optimizers
 
     def __dataloader(self, train):
         # init data generators
@@ -189,6 +201,12 @@ class ExampleModel(LightningModule):
 
     @staticmethod
     def add_model_specific_args(parent_parser, root_dir):
+        """
+        Parameters you define here will be available to your model through self.hparams
+        :param parent_parser:
+        :param root_dir:
+        :return:
+        """
         parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
 
         # param overwrites
@@ -209,3 +227,5 @@ class ExampleModel(LightningModule):
         parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
         parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
         return parser
+
+```

docs/Trainer/Training Loop.md

@@ -38,6 +38,16 @@ Use this to turn off early stopping and run training to the [max_epoch](#force-t
 trainer = Trainer(enable_early_stop=True)
 ```
 
+---
+#### Gradient Clipping 
+Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
+``` {.python}
+# DEFAULT (ie: don't clip)
+trainer = Trainer(gradient_clip=0)
+```
+
+
+
 ---
 #### Inspect gradient norms
 Looking at grad norms can help you figure out where training might be going wrong.

docs/index.md

@@ -62,6 +62,7 @@
 - [Anneal Learning rate](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#anneal-learning-rate)
 - [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
 - [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
+- [Gradient Clipping: DOC TODO](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
 - [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/Pytorch-Lightning/LightningModule/#configure_optimizers)
 - [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
 

docs/source/examples/__init__.py

@@ -1 +0,0 @@
-from .example_model import ExampleModel

docs/source/examples/basic_trainer.py

@@ -1,74 +0,0 @@
-import os
-import sys
-
-from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
-from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
-from docs.source.examples.example_model import ExampleModel
-
-
-def main(hparams):
-    """
-    Main training routine specific for this project
-    :param hparams:
-    :return:
-    """
-    # init experiment
-    exp = Experiment(
-        name=hparams.tt_name,
-        debug=hparams.debug,
-        save_dir=hparams.tt_save_path,
-        version=hparams.hpc_exp_number,
-        autosave=False,
-        description=hparams.tt_description
-    )
-
-    exp.argparse(hparams)
-    exp.save()
-
-    # build model
-    model = ExampleModel(hparams)
-
-    # callbacks
-    early_stop = EarlyStopping(
-        monitor='val_acc',
-        patience=3,
-        mode='min',
-        verbose=True,
-    )
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-    checkpoint = ModelCheckpoint(
-        filepath=model_save_path,
-        save_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor='val_acc',
-        mode='min'
-    )
-
-    # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-    )
-
-    # train model
-    trainer.fit(model)
-
-
-if __name__ == '__main__':
-
-    # use default args given by lightning
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
-    parent_parser = HyperOptArgumentParser(strategy='random_search', add_help=False)
-    add_default_args(parent_parser, root_dir)
-
-    # allow model to overwrite or extend args
-    parser = ExampleModel.add_model_specific_args(parent_parser)
-    hyperparams = parser.parse_args()
-
-    # train model
-    main(hyperparams)

docs/source/examples/fully_featured_trainer.py

@@ -1,210 +0,0 @@
-import os
-import sys
-import numpy as np
-from time import sleep
-import torch
-
-from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
-
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
-
-# ---------------------
-# DEFINE MODEL HERE
-# ---------------------
-from docs.source.examples.example_model import ExampleModel
-# ---------------------
-
-AVAILABLE_MODELS = {
-    'model_template': ExampleModel
-}
-
-
-"""
-Allows training by using command line arguments
-Run by:
-# TYPE YOUR RUN COMMAND HERE
-"""
-
-
-def main_local(hparams):
-    main(hparams, None, None)
-
-
-def main(hparams, cluster, results_dict):
-    """
-    Main training routine specific for this project
-    :param hparams:
-    :return:
-    """
-    on_gpu = hparams.gpus is not None and torch.cuda.is_available()
-
-    device = 'cuda' if on_gpu else 'cpu'
-    hparams.__setattr__('device', device)
-    hparams.__setattr__('on_gpu', on_gpu)
-    hparams.__setattr__('nb_gpus', torch.cuda.device_count())
-    hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
-
-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
-    # init experiment
-    log_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(
-        name='test_tube_exp',
-        debug=True,
-        save_dir=log_dir,
-        version=0,
-        autosave=False,
-        description='test demo'
-    )
-
-    exp.argparse(hparams)
-    exp.save()
-
-    # build model
-    print('loading model...')
-    model = TRAINING_MODEL(hparams)
-    print('model built')
-
-    # callbacks
-    early_stop = EarlyStopping(
-        monitor=hparams.early_stop_metric,
-        patience=hparams.early_stop_patience,
-        verbose=True,
-        mode=hparams.early_stop_mode
-    )
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-    checkpoint = ModelCheckpoint(
-        filepath=model_save_path,
-        save_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor=hparams.model_save_monitor_value,
-        mode=hparams.model_save_monitor_mode
-    )
-
-    # gpus are ; separated for inside a node and , within nodes
-    gpu_list = None
-    if hparams.gpus is not None:
-        gpu_list = [int(x) for x in hparams.gpus.split(';')]
-
-    # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        cluster=cluster,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-        gpus=gpu_list
-    )
-
-    # train model
-    trainer.fit(model)
-
-
-def get_default_parser(strategy, root_dir):
-
-    possible_model_names = list(AVAILABLE_MODELS.keys())
-    parser = HyperOptArgumentParser(strategy=strategy, add_help=False)
-    add_default_args(parser, root_dir, possible_model_names=possible_model_names, rand_seed=SEED)
-    return parser
-
-
-def get_model_name(args):
-    for i, arg in enumerate(args):
-        if 'model_name' in arg:
-            return args[i+1]
-
-
-def optimize_on_cluster(hyperparams):
-    # enable cluster training
-    cluster = SlurmCluster(
-        hyperparam_optimizer=hyperparams,
-        log_path=hyperparams.tt_save_path,
-        test_tube_exp_name=hyperparams.tt_name
-    )
-
-    # email for cluster coms
-    cluster.notify_job_status(email='add_email_here', on_done=True, on_fail=True)
-
-    # configure cluster
-    cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus
-    cluster.job_time = '48:00:00'
-    cluster.gpu_type = '1080ti'
-    cluster.memory_mb_per_node = 48000
-
-    # any modules for code to run in env
-    cluster.add_command('source activate pytorch_lightning')
-
-    # name of exp
-    job_display_name = hyperparams.tt_name.split('_')[0]
-    job_display_name = job_display_name[0:3]
-
-    # run hopt
-    print('submitting jobs...')
-    cluster.optimize_parallel_cluster_gpu(
-        main,
-        nb_trials=hyperparams.nb_hopt_trials,
-        job_name=job_display_name
-    )
-
-
-if __name__ == '__main__':
-
-    model_name = get_model_name(sys.argv)
-    if model_name is None:
-        model_name = 'model_template'
-
-    # use default args
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    parent_parser = get_default_parser(strategy='random_search', root_dir=root_dir)
-
-    # allow model to overwrite or extend args
-    TRAINING_MODEL = AVAILABLE_MODELS[model_name]
-    parser = TRAINING_MODEL.add_model_specific_args(parent_parser, root_dir)
-    hyperparams = parser.parse_args()
-
-    # format GPU layout
-    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-
-    # ---------------------
-    # RUN TRAINING
-    # ---------------------
-
-    # cluster and CPU
-    if hyperparams.on_cluster:
-        # run on HPC cluster
-        print('RUNNING ON SLURM CLUSTER')
-        gpu_ids = hyperparams.gpus.split(';')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        optimize_on_cluster(hyperparams)
-
-    elif hyperparams.gpus is None:
-        # run on cpu
-        print('RUNNING ON CPU')
-        main(hyperparams, None, None)
-
-    # single or multiple GPUs on same machine
-    gpu_ids = hyperparams.gpus.split(';')
-    if hyperparams.interactive:
-        # run on 1 gpu
-        print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {gpu_ids}')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        main(hyperparams, None, None)
-
-    else:
-        # multiple GPUs on same machine
-        print(f'RUNNING MULTI GPU. GPU ids: {gpu_ids}')
-        hyperparams.optimize_parallel_gpu(
-            main_local,
-            gpu_ids=gpu_ids,
-            nb_trials=hyperparams.nb_hopt_trials,
-            nb_workers=len(gpu_ids)
-        )

pytorch_lightning/models/trainer.py

@@ -33,6 +33,7 @@ class Trainer(TrainerIO):
     def __init__(self,
                  experiment,
                  checkpoint_callback, early_stop_callback,
+                 gradient_clip=0,
                  cluster=None,
                  process_position=0,
                  current_gpu_name=0,
@@ -53,6 +54,7 @@ class Trainer(TrainerIO):
                  nb_sanity_val_steps=5):
 
         # Transfer params
+        self.gradient_clip = gradient_clip
         self.check_val_every_n_epoch = check_val_every_n_epoch
         self.enable_early_stop = enable_early_stop
         self.track_grad_norm = track_grad_norm
@@ -441,6 +443,11 @@ class Trainer(TrainerIO):
         # gradient update with accumulated gradients
         if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
 
+            # clip gradients
+            if self.gradient_clip > 0:
+                model = self.model.module if self.data_parallel else self.model
+                torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
+
             # update gradients across all optimizers
             for optimizer in self.optimizers:
                 optimizer.step()