release v0.2.4.1

add a hook for on_tng_metrics so that users get access to the grad_no…

.gitignore

@@ -8,6 +8,7 @@ datasets/
 model_weights/
 app/models/
 pip-wheel-metadata/
+test_tube_exp/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -36,9 +36,26 @@ To use lightning do 2 things:
 2. [Define a LightningModel](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py).     
 
 ## What does lightning control for me?
-Everything! Except the following three things:
+Everything!    
+Except for these 6 core functions which you define:    
 
-**What happens in the training loop**    
+```{.python}
+# what to do in the training loop
+def training_step(self, data_batch, batch_nb):
+
+# what to do in the validation loop
+def validation_step(self, data_batch, batch_nb):
+
+# how to aggregate validation_step outputs
+def validation_end(self, outputs):
+
+# and your dataloaders
+def tng_dataloader():
+def val_dataloader():
+def test_dataloader():
+```
+
+**Could be as complex as seq-2-seq + attention**    
 
 ```python
 # define what happens for training here
@@ -46,25 +63,39 @@ def training_step(self, data_batch, batch_nb):
     x, y = data_batch
     
     # define your own forward and loss calculation
-    out = self.forward(x)
+    hidden_states = self.encoder(x)
-    loss = my_loss(out, y)
+     
+    # even as complex as a seq-2seq + attn model
+    # (this is just a toy, non-working example to illustrate)
+    start_token = '<SOS>'
+    last_hidden = torch.zeros(...)
+    loss = 0
+    for step in range(max_seq_len):
+        attn_context = self.attention_nn(hidden_states, start_token)
+        pred = self.decoder(start_token, attn_context, last_hidden) 
+        last_hidden = pred
+        pred = self.predict_nn(pred)
+        loss += self.loss(last_hidden, y[step])
+        
+    #toy example as well
+    loss = loss / max_seq_len
     return {'loss': loss} 
 ```
 
-**What happens in the validation loop**      
+**Or as basic as CNN image classification**      
 
 ```python
 # define what happens for validation here
 def validation_step(self, data_batch, batch_nb):    
     x, y = data_batch
     
-    # define your own forward and loss calculation
+    # or as basic as a CNN classification
     out = self.forward(x)
     loss = my_loss(out, y)
     return {'loss': loss} 
 ```
 
-**And what to do with the output of all validation batches**    
+**And you also decide how to collate the output of all validation steps**    
 
 ```python
 def validation_end(self, outputs):
@@ -84,8 +115,40 @@ def validation_end(self, outputs):
     tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
     return tqdm_dic
 ```
+   
+## Tensorboard    
+Lightning is fully integrated with tensorboard.   
 
-## Lightning gives you options to control the following:
+<p align="center">
+  <a href="https://williamfalcon.github.io/pytorch-lightning/">
+    <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/tf_loss.png" width="900px">
+  </a>
+</p>
+
+Lightning also adds a text column with all the hyperparameters for this experiment.      
+
+<p align="center">
+  <a href="https://williamfalcon.github.io/pytorch-lightning/">
+        <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/tf_tags.png" width="900px">
+  </a>
+</p>
+
+Simply note the path you set for the Experiment    
+``` {.python}   
+from test_tube import Experiment
+from pytorch-lightning import  Trainer
+
+exp = Experiment(save_dir='/some/path')
+trainer = Trainer(experiment=exp)
+...
+```   
+
+And run tensorboard from that dir   
+```bash
+tensorboard --logdir /some/path     
+```    
+
+## Lightning automatically automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
 
 ###### Checkpointing    
 

docs/LightningModule/properties.md

@@ -8,12 +8,21 @@ The current epoch
 #### dtype    
 Current dtype    
 
+---
+#### experiment    
+An instance of test-tube Experiment which you can use to log anything for tensorboarX.   
+```{.python}
+self.experiment.add_embedding(...)   
+self.experiment.log({'val_loss': 0.9})   
+self.experiment.add_scalars(...)   
+```
+
 --- 
-#### global_step
+#### global_step   
 Total training batches seen across all epochs   
 
 --- 
-#### gradient_clip
+#### gradient_clip   
 The current gradient clip value    
 
 ---
@@ -21,5 +30,11 @@ The current gradient clip value
 True if your model is currently running on GPUs. Useful to set flags around the LightningModule for different CPU vs GPU behavior.    
 
 ---
-#### Trainer
+#### trainer    
 Last resort access to any state the trainer has. Changing certain properties here could affect your training run.
+```{.python}   
+self.trainer.optimizers
+self.trainer.current_epoch
+...   
+```   
+

examples/new_project_templates/__init__.py

@@ -0,0 +1 @@
+from .lightning_module_template import LightningTemplateModel

examples/new_project_templates/lightning_module_template.py

@@ -7,6 +7,8 @@ import torch
 import torch.nn.functional as F
 from test_tube import HyperOptArgumentParser
 from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
 
 from pytorch_lightning.root_module.root_module import LightningModule
 
@@ -154,13 +156,25 @@ class LightningTemplateModel(LightningModule):
     def __dataloader(self, train):
         # init data generators
         transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-
         dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
 
-        loader = torch.utils.data.DataLoader(
+        # when using multi-node we need to add the datasampler
+        train_sampler = None
+        batch_size = self.hparams.batch_size
+
+        try:
+            if self.on_gpu:
+                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+                batch_size = batch_size // self.trainer.world_size  # scale batch size
+        except Exception as e:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
             dataset=dataset,
-            batch_size=self.hparams.batch_size,
+            batch_size=batch_size,
-            shuffle=True
+            shuffle=should_shuffle,
+            sampler=train_sampler
         )
 
         return loader
@@ -210,16 +224,19 @@ class LightningTemplateModel(LightningModule):
 
         # network params
         parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
-        parser.add_argument('--in_features', default=28*28)
+        parser.add_argument('--in_features', default=28*28, type=int)
-        parser.add_argument('--out_features', default=10)
+        parser.add_argument('--out_features', default=10, type=int)
-        parser.add_argument('--hidden_dim', default=50000) # use 500 for CPU, 50000 for GPU to see speed difference
+        parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference
 
         # data
         parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
 
         # training params (opt)
-        parser.opt_list('--learning_rate', default=0.001, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
+        parser.opt_list('--learning_rate', default=0.001*8, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
                         tunable=False)
-        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)
+
+        # if using 2 nodes with 4 gpus each the batch size here (256) will be 256 / (2*8) = 16 per gpu
+        parser.opt_list('--batch_size', default=256*8, type=int, options=[32, 64, 128, 256], tunable=False,
+                        help='batch size will be divided over all the gpus being used across all nodes')
         return parser

examples/new_project_templates/multi_node_cluster_template.py

@@ -0,0 +1,172 @@
+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 lightning_module_template import LightningTemplateModel
+# ---------------------
+
+"""
+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:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+    # when using grid search, it's possible for all models to start at once
+    # and use the same test tube experiment version
+    relative_node_id = int(os.environ['SLURM_NODEID'])
+    sleep(relative_node_id + 1)
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        cluster=cluster,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+        nb_gpu_nodes=hyperparams.nb_gpu_nodes
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+def optimize_on_cluster(hyperparams):
+    # enable cluster training
+    # log all scripts to the test tube folder
+    cluster = SlurmCluster(
+        hyperparam_optimizer=hyperparams,
+        log_path=hyperparams.slurm_log_path,
+    )
+
+    # 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.per_experiment_nb_nodes = hyperparams.nb_gpu_nodes
+    cluster.job_time = '2:00:00'
+    cluster.gpu_type = 'volta'
+    cluster.memory_mb_per_node = 0
+
+    # any modules for code to run in env
+    cluster.add_command('source activate lightning')
+
+    # run only on 32GB voltas
+    cluster.add_slurm_cmd(cmd='constraint', value='volta32gb', comment='use 32gb gpus')
+    cluster.add_slurm_cmd(cmd='partition', value=hyperparams.gpu_partition, comment='use 32gb gpus')
+
+    # run hopt
+    # creates and submits jobs to slurm
+    cluster.optimize_parallel_cluster_gpu(
+        main,
+        nb_trials=hyperparams.nb_hopt_trials,
+        job_name=hyperparams.experiment_name
+    )
+
+
+if __name__ == '__main__':
+
+    # use default args
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+    slurm_out_dir = os.path.join(demo_log_dir, 'slurm_scripts')
+
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # cluster args not defined inside the model
+    parent_parser.add_argument('--gpu_partition', type=str, help='consult your cluster manual')
+
+    # TODO: make 1 param
+    parent_parser.add_argument('--per_experiment_nb_gpus', type=int, help='how many gpus to use in a node')
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node')
+
+    parent_parser.add_argument('--nb_gpu_nodes', type=int, default=1, help='how many nodes to use in a cluster')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--slurm_log_path', type=str, default=slurm_out_dir, help='where to save slurm meta')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+    parent_parser.add_argument('--nb_hopt_trials', type=int, default=1, help='how many grid search trials to run')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print('RUNNING ON SLURM CLUSTER')
+    optimize_on_cluster(hyperparams)

examples/new_project_templates/single_cpu_template.py

@@ -0,0 +1,110 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+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)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING ON CPU')
+    main(hyperparams)

examples/new_project_templates/single_gpu_node_16bit_template.py

@@ -0,0 +1,113 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+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)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+        use_amp=True
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)

examples/new_project_templates/single_gpu_node_template.py

@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+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)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    print('loading model...')
+    model = LightningTemplateModel(hparams)
+    print('model built')
+
+    # ------------------------
+    # 2 INIT TEST TUBE EXP
+    # ------------------------
+
+    # init experiment
+    exp = Experiment(
+        name=hyperparams.experiment_name,
+        save_dir=hyperparams.test_tube_save_path,
+        autosave=False,
+        description='test demo'
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    # ------------------------
+    # 3 DEFINE CALLBACKS
+    # ------------------------
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+    early_stop = EarlyStopping(
+        monitor='val_acc',
+        patience=3,
+        verbose=True,
+        mode='max'
+    )
+
+    checkpoint = ModelCheckpoint(
+        filepath=model_save_path,
+        save_best_only=True,
+        verbose=True,
+        monitor='val_loss',
+        mode='min'
+    )
+
+    # ------------------------
+    # 4 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        early_stop_callback=early_stop,
+        gpus=hparams.gpus,
+    )
+
+    # ------------------------
+    # 5 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # dirs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+    checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+    test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+    # although we user hyperOptParser, we are using it only as argparse right now
+    parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+    # gpu args
+    parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+    parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+    parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+    parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+    # allow model to overwrite or extend args
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    # run on HPC cluster
+    print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+    main(hyperparams)

examples/new_project_templates/trainer_gpu_cluster_template.py

@@ -1,207 +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 lightning_module_template import LightningTemplateModel
-# ---------------------
-
-AVAILABLE_MODELS = {
-    'model_template': LightningTemplateModel
-}
-
-
-"""
-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',
-        save_dir=log_dir,
-        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_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/__init__.py

@@ -1 +1,2 @@
-from .models import Trainer
+from .models import Trainer
+from .root_module.root_module import LightningModule

pytorch_lightning/callbacks/pt_callbacks.py

@@ -1,6 +1,6 @@
 import numpy as np
 import os, shutil
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 
 
 class Callback(object):
@@ -34,7 +34,7 @@ class Callback(object):
         self.params = params
 
     def set_model(self, model):
-        if type(model) is LightningDataParallel:
+        if type(model) is LightningDistributedDataParallel:
             model = model.module
         self.model = model
 

pytorch_lightning/models/trainer.py

@@ -1,12 +1,24 @@
-import torch
+"""
-import tqdm
+The trainer handles all the logic for running a val loop, training loop, distributing, etc...
-import numpy as np
+"""
-from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from time import sleep
+import subprocess
 import traceback
-from pytorch_lightning.root_module.model_saving import TrainerIO
+import warnings
+import os
+
+import torch
+from torch.utils.data.distributed import DistributedSampler
 from torch.optim.lr_scheduler import MultiStepLR
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
+import torch.multiprocessing as mp
-import pdb
+import torch.distributed as dist
+import numpy as np
+import tqdm
+
+from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module.model_saving import TrainerIO
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
+
 
 try:
     from apex import amp
@@ -15,28 +27,17 @@ except ModuleNotFoundError:
     APEX_AVAILABLE = False
 
 
-def reduce_distributed_output(output, nb_gpus):
-    for k, v in output.items():
-        # recurse on nested dics
-        if isinstance(output[k], dict):
-            output[k] = reduce_distributed_output(output[k], nb_gpus)
-
-        # reduce only metrics that have the same nb of gpus
-        elif output[k].size(0) == nb_gpus:
-            reduced = torch.mean(output[k])
-            output[k] = reduced
-    return output
-
-
 class Trainer(TrainerIO):
 
     def __init__(self,
                  experiment,
-                 checkpoint_callback, early_stop_callback,
+                 early_stop_callback=None,
+                 checkpoint_callback=None,
                  gradient_clip=0,
                  cluster=None,
                  process_position=0,
                  current_gpu_name=0,
+                 nb_gpu_nodes=1,
                  gpus=None,
                  progress_bar=True,
                  overfit_pct=0.0,
@@ -44,19 +45,22 @@ class Trainer(TrainerIO):
                  check_val_every_n_epoch=1,
                  fast_dev_run=False,
                  accumulate_grad_batches=1,
-                 enable_early_stop=True, max_nb_epochs=1000, min_nb_epochs=1,
+                 max_nb_epochs=1000, min_nb_epochs=1,
-                 train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0, val_check_interval=0.95,
+                 train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0,
+                 val_check_interval=0.95,
                  log_save_interval=100, add_log_row_interval=10,
                  lr_scheduler_milestones=None,
                  use_amp=False,
                  print_nan_grads=False,
+                 print_weights_summary=True,
                  amp_level='O2',
                  nb_sanity_val_steps=5):
 
         # Transfer params
+        self.nb_gpu_nodes = nb_gpu_nodes
         self.gradient_clip = gradient_clip
         self.check_val_every_n_epoch = check_val_every_n_epoch
-        self.enable_early_stop = enable_early_stop
+        self.enable_early_stop = early_stop_callback is not None
         self.track_grad_norm = track_grad_norm
         self.fast_dev_run = fast_dev_run
         self.on_gpu = gpus is not None and torch.cuda.is_available()
@@ -66,8 +70,12 @@ class Trainer(TrainerIO):
         self.cluster = cluster
         self.process_position = process_position
         self.current_gpu_name = current_gpu_name
+        self.print_weights_summary = print_weights_summary
         self.checkpoint_callback = checkpoint_callback
-        self.checkpoint_callback.save_function = self.save_checkpoint
+
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.save_function = self.save_checkpoint
+
         self.early_stop = early_stop_callback
         self.model = None
         self.max_nb_epochs = max_nb_epochs
@@ -79,8 +87,26 @@ class Trainer(TrainerIO):
         self.lr_schedulers = []
         self.amp_level = amp_level
         self.print_nan_grads = print_nan_grads
-        self.data_parallel_device_ids = gpus
+        self.data_parallel_device_ids = None
-        self.data_parallel = gpus is not None and len(gpus) > 0
+        self.world_size = 1
+
+        # gpus come in as a string.
+        # if gpus = -1 then use all available devices
+        # otherwise, split the string using commas
+        if gpus is not None:
+            if gpus == '-1':
+                self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
+            else:
+                self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
+
+            # set the correct cuda visible devices (using pci order)
+            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+            os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
+
+        self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 1
+
+        # process info
+        self.proc_rank = 0
 
         # training state
         self.optimizers = None
@@ -103,11 +129,21 @@ class Trainer(TrainerIO):
         self.__determine_data_use_amount(train_percent_check, val_percent_check, test_percent_check, overfit_pct)
         print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
 
-        # apex test
+        # 16 bit mixed precision training using apex
         self.use_amp = use_amp and APEX_AVAILABLE
         if self.use_amp:
             print('using 16bit precision')
 
+        if use_amp and not APEX_AVAILABLE:
+            msg = '''
+            You set use_amp=True but do not have apex installed.
+            Install apex first using this guide and rerun with use_amp=True: 
+            https://github.com/NVIDIA/apex#linux
+            
+            this run will NOT use 16 bit precision
+            '''
+            warnings.warn(msg)
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -120,8 +156,12 @@ class Trainer(TrainerIO):
             self.val_percent_check = overfit_pct
             self.test_percent_check = overfit_pct
 
+    def __get_model(self):
+        return self.model.module if self.data_parallel else self.model
+
     def __is_function_implemented(self, f_name):
-        f_op = getattr(self.model, f_name, None)
+        model = self.__get_model()
+        f_op = getattr(model, f_name, None)
         return callable(f_op)
 
     @property
@@ -139,7 +179,7 @@ class Trainer(TrainerIO):
 
         return tqdm_dic
 
-    def __layout_bookeeping(self, model):
+    def __layout_bookeeping(self):
         # training bookeeping
         self.total_batch_nb = 0
         self.running_loss = []
@@ -148,17 +188,17 @@ class Trainer(TrainerIO):
         self.tqdm_metrics = {}
 
         # determine number of training batches
-        self.nb_tng_batches = model.nb_batches(self.tng_dataloader)
+        self.nb_tng_batches = len(self.tng_dataloader)
         self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
 
         # determine number of validation batches
-        self.nb_val_batches = model.nb_batches(self.val_dataloader)
+        self.nb_val_batches = len(self.val_dataloader)
         self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
         self.nb_val_batches = max(1, self.nb_val_batches)
         self.nb_val_batches = self.nb_val_batches
 
         # determine number of test batches
-        self.nb_test_batches = model.nb_batches(self.test_dataloader)
+        self.nb_test_batches = len(self.test_dataloader)
         self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
 
         # determine when to check validation
@@ -179,12 +219,9 @@ class Trainer(TrainerIO):
         :param max_batches: Scalar
         :return:
         """
-        print('validating...')
-
         # enable eval mode
         model.zero_grad()
         model.eval()
-        model.from_lightning = True
 
         # disable gradients to save memory
         torch.set_grad_enabled(False)
@@ -207,7 +244,6 @@ class Trainer(TrainerIO):
             # -----------------
             if self.data_parallel:
                 output = model(data_batch, batch_i)
-                output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
             else:
                 output = model.validation_step(data_batch, batch_i)
 
@@ -241,32 +277,168 @@ class Trainer(TrainerIO):
         self.test_dataloader = model.test_dataloader
         self.val_dataloader = model.val_dataloader
 
+        if self.data_parallel and not isinstance(self.tng_dataloader.sampler, DistributedSampler):
+            msg = '''
+            when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
+            
+            ie: this:
+            dataset = myDataset()
+            dataloader = Dataloader(dataset)
+            
+            becomes:
+            dataset = myDataset()
+            dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+            dataloader = Dataloader(dataset, sampler=dist_sampler)
+            '''
+            raise Exception(msg)
+
     # -----------------------------
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
 
-        # give model convenience properties
+        # when using gpus, first thing we do is spawn a new process between each worker
-        model.trainer = self
+        #  multi-gpu and multi-nodes
-        model.experiment = self.experiment
+        if self.data_parallel:
+            self.experiment = self.experiment.get_meta_copy()
+            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
 
-        # transfer data loaders from model
+        # treat 1 gpu as a different case to avoid nccl bugs
-        self.__get_dataloaders(model)
+        elif self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) == 1:
+            self.single_gpu_train(model)
 
-        # init training constants
+        else:
-        self.__layout_bookeeping(model)
+            # CHOOSE OPTIMIZER
+            # filter out the weights that were done on gpu so we can load on good old cpus
+            self.optimizers = model.configure_optimizers()
+
+            # run through amp wrapper
+            if self.use_amp:
+                # An example
+                model, optimizers = amp.initialize(
+                    model, self.optimizers, opt_level=self.amp_level,
+                )
+                self.optimizers = optimizers
+
+            self.__run_pretrain_routine(model)
+
+    def single_gpu_train(self, model):
+        # torch.cuda.set_device(0)
+        model.cuda(0)
 
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
 
+        # run through amp wrapper
         if self.use_amp:
             # An example
-            model, optimizer = amp.initialize(
+            model, optimizers = amp.initialize(
-                model, self.optimizers[0], opt_level=self.amp_level,
+                model, self.optimizers, opt_level=self.amp_level,
             )
-            self.optimizers[0] = optimizer
+            self.optimizers = optimizers
-            model.trainer = self
+
+        self.__run_pretrain_routine(model)
+
+    def dp_train(self, gpu_nb, model):
+        """
+        Entry point into a DP thread
+        :param gpu_nb:
+        :param model:
+        :param cluster_obj:
+        :return:
+        """
+        # node rank using relative slurm id
+        # otherwise default to node rank 0
+        try:
+            node_rank = int(os.environ['SLURM_NODEID'])
+        except KeyError as e:
+            node_rank = 0
+
+        # recover original exp before went into process
+        # init in write mode only on proc 0
+        self.experiment.debug = self.proc_rank > 0
+        self.experiment = self.experiment.get_non_ddp_exp()
+
+        # show progbar only on prog_rank 0
+        self.prog_bar = self.prog_bar and node_rank == 0 and gpu_nb == 0
+
+        # determine which process we are and world size
+        self.proc_rank = node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
+
+        # set up server using proc 0's ip address
+        # try to init for 20 times at max in case ports are taken
+        # where to store ip_table
+        self.__init_tcp_connection()
+
+        # CHOOSE OPTIMIZER
+        # filter out the weights that were done on gpu so we can load on good old cpus
+        self.optimizers = model.configure_optimizers()
+
+        # MODEL
+        # copy model to each gpu
+        torch.cuda.set_device(gpu_nb)
+        model.cuda(gpu_nb)
+
+        # AMP
+        # run through amp wrapper before going to distributed DP
+        if self.use_amp:
+            # An example
+            model, optimizers = amp.initialize(
+                model, self.optimizers, opt_level=self.amp_level,
+            )
+            self.optimizers = optimizers
+
+        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb])
+
+        # continue training routine
+        self.__run_pretrain_routine(model)
+
+    def __init_tcp_connection(self):
+        """
+        Connect all procs in the world using the env:// init
+        Use the first node as the root address
+        :param port:
+        :param tries:
+        :return:
+        """
+        try:
+            port = os.environ['MASTER_PORT']
+        except Exception as e:
+            port = 12910
+            os.environ['MASTER_PORT'] = f'{port}'
+
+        try:
+            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        except Exception as e:
+            root_node = '127.0.0.2'
+
+        os.environ['MASTER_ADDR'] = root_node
+
+        sleep(self.proc_rank*0.5)
+        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+
+    def __run_pretrain_routine(self, model):
+        """
+        Sanity check a few things before starting actual training
+        :param model:
+        :return:
+        """
+        ref_model = model
+        if self.data_parallel:
+            ref_model = model.module
+
+        ref_model.trainer = self
+
+        # set local properties on the model
+        ref_model.on_gpu = self.on_gpu
+
+        # transfer data loaders from model
+        self.__get_dataloaders(ref_model)
+
+        # init training constants
+        self.__layout_bookeeping()
 
         # add lr schedulers
         if self.lr_scheduler_milestones is not None:
@@ -275,17 +447,19 @@ class Trainer(TrainerIO):
                 self.lr_schedulers.append(scheduler)
 
         # print model summary
-        model.summarize()
+        if self.proc_rank == 0 and self.print_weights_summary:
+            ref_model.summarize()
 
-        # put on gpu if needed
+        # give model convenience properties
-        if self.on_gpu:
+        ref_model.trainer = self
-            model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+        ref_model.experiment = self.experiment
 
         # run tiny validation to make sure program won't crash during val
         _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
 
         # save exp to get started
-        self.experiment.save()
+        if self.proc_rank == 0:
+            self.experiment.save()
 
         # enable cluster checkpointing
         if self.cluster is not None:
@@ -304,12 +478,12 @@ class Trainer(TrainerIO):
             for lr_scheduler in self.lr_schedulers:
                 lr_scheduler.step()
 
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             model.current_epoch = epoch_nb
 
             # hook
             if self.__is_function_implemented('on_epoch_start'):
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 model.on_epoch_start()
 
             self.current_epoch = epoch_nb
@@ -324,7 +498,7 @@ class Trainer(TrainerIO):
                 self.batch_nb = batch_nb
                 self.global_step += 1
 
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 model.global_step = self.global_step
 
                 # stop when the flag is changed or we've gone past the amount requested in the batches
@@ -348,17 +522,16 @@ class Trainer(TrainerIO):
 
                 # when batch should be saved
                 if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
-                    self.experiment.save()
+                    if self.proc_rank == 0:
+                        self.experiment.save()
 
                 # when metrics should be logged
                 if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
-                    if self.data_parallel:
+                    model = self.__get_model()
-                        metrics = self.model.module.update_tng_log_metrics(self.__tng_tqdm_dic)
+                    metrics = model.update_tng_log_metrics(self.__tng_tqdm_dic)
-                    else:
-                        metrics = self.model.update_tng_log_metrics(self.__tng_tqdm_dic)
 
                     # add gpu memory
                     if self.on_gpu:
@@ -367,19 +540,22 @@ class Trainer(TrainerIO):
 
                     # add norms
                     if self.track_grad_norm > 0:
-                        model = self.model.module if self.data_parallel else self.model
+                        model = self.__get_model()
                         grad_norm_dic = model.grad_norm(self.track_grad_norm)
-
                         metrics.update(grad_norm_dic)
 
+                    if self.__is_function_implemented('on_tng_metrics'):
+                        model.on_tng_metrics(metrics)
+
                     # log metrics
                     scalar_metrics = self.__metrics_to_scalars(metrics, blacklist=self.__log_vals_blacklist())
-                    self.experiment.log(scalar_metrics, global_step=self.global_step)
+                    if self.proc_rank == 0:
-                    self.experiment.save()
+                        self.experiment.log(scalar_metrics, global_step=self.global_step)
+                        self.experiment.save()
 
                 # hook
                 if self.__is_function_implemented('on_batch_end'):
-                    model = self.model.module if self.data_parallel else self.model
+                    model = self.__get_model()
                     model.on_batch_end()
 
                 # end epoch early
@@ -388,13 +564,13 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_epoch_end'):
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 model.on_epoch_end()
 
             # early stopping
-            if self.enable_early_stop:
+            met_min_epochs = epoch_nb > self.min_nb_epochs
+            if self.enable_early_stop and met_min_epochs:
                 should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, logs=self.__tng_tqdm_dic)
-                met_min_epochs = epoch_nb > self.min_nb_epochs
 
                 # stop training
                 stop = should_stop and met_min_epochs
@@ -426,7 +602,7 @@ class Trainer(TrainerIO):
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             response = model.on_batch_start(data_batch)
 
             if response == -1:
@@ -439,17 +615,27 @@ class Trainer(TrainerIO):
         # return a scalar value and a dic with tqdm metrics
         if self.data_parallel:
             output = self.model(data_batch, batch_nb)
-            output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
         else:
             output = self.model.training_step(data_batch, batch_nb)
 
-        model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        try:
-        loss = output['loss']
+            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        except Exception as e:
+            model_specific_tqdm_metrics_dic = {}
+
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        try:
+            loss = output['loss']
+        except Exception as e:
+            if type(output) is torch.Tensor:
+                loss = output
 
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
         if self.use_amp:
+            # scale loss when using amp
             for optimizer in self.optimizers:
                 with amp.scale_loss(loss, optimizer) as scaled_loss:
                     scaled_loss.backward()
@@ -457,10 +643,11 @@ class Trainer(TrainerIO):
             loss.backward()
 
         if self.print_nan_grads:
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             for param in model.parameters():
                 print(param.grad.float().sum())
 
+        # avoid memory leaks
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients
@@ -468,7 +655,7 @@ class Trainer(TrainerIO):
 
             # clip gradients
             if self.gradient_clip > 0:
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
 
             # update gradients across all optimizers
@@ -494,7 +681,8 @@ class Trainer(TrainerIO):
 
         # activate batch end hook
         if self.__is_function_implemented('on_batch_end'):
-            self.model.on_batch_end()
+            model = self.__get_model()
+            model.on_batch_end()
 
         return 0
 
@@ -509,7 +697,8 @@ class Trainer(TrainerIO):
         try:
             # hook
             if self.__is_function_implemented('on_pre_performance_check'):
-                self.model.on_pre_performance_check()
+                model = self.__get_model()
+                model.on_pre_performance_check()
 
             # use full val set on end of epoch
             # use a small portion otherwise
@@ -523,7 +712,8 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_post_performance_check'):
-                self.model.on_post_performance_check()
+                model = self.__get_model()
+                model.on_post_performance_check()
 
         except Exception as e:
             print(e)
@@ -535,5 +725,6 @@ class Trainer(TrainerIO):
             self.prog_bar.set_postfix(**tqdm_metrics)
 
         # model checkpointing
-        print('save callback...')
+        if self.proc_rank == 0 and self.checkpoint_callback:
-        self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
+            print('save callback...')
+            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -1,4 +1,6 @@
 from torch.nn import DataParallel
+from torch.nn.parallel import DistributedDataParallel
+import itertools
 
 import threading
 import torch
@@ -6,6 +8,20 @@ from torch.cuda._utils import _get_device_index
 import pdb
 
 
+def _find_tensors(obj):
+    r"""
+    Recursively find all tensors contained in the specified object.
+    """
+    if isinstance(obj, torch.Tensor):
+        return [obj]
+    if isinstance(obj, (list, tuple)):
+        return itertools.chain(*map(_find_tensors, obj))
+    if isinstance(obj, dict):
+        return itertools.chain(*map(_find_tensors, obj.values()))
+    return []
+
+
+
 def get_a_var(obj):
     if isinstance(obj, torch.Tensor):
         return obj
@@ -30,6 +46,49 @@ class LightningDataParallel(DataParallel):
         return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
 
 
+class LightningDistributedDataParallel(DistributedDataParallel):
+    """
+    Override the forward call in lightning so it goes to training and validation step respectively
+    """
+
+    def parallel_apply(self, replicas, inputs, kwargs):
+        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
+
+    def forward(self, *inputs, **kwargs):
+        self._sync_params()
+        if self.device_ids:
+            inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
+            if len(self.device_ids) == 1:
+                # --------------
+                # LIGHTNING MOD
+                # --------------
+                # normal
+                # output = self.module(*inputs[0], **kwargs[0])
+
+                # lightning
+                if self.module.training:
+                    output = self.module.training_step(*inputs[0], **kwargs[0])
+                else:
+                    output = self.module.validation_step(*inputs[0], **kwargs[0])
+            else:
+                outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
+                output = self.gather(outputs, self.output_device)
+        else:
+            output = self.module(*inputs, **kwargs)
+
+        if torch.is_grad_enabled():
+            # We'll return the output object verbatim since it is a freeform
+            # object. We need to find any tensors in this object, though,
+            # because we need to figure out which parameters were used during
+            # this forward pass, to ensure we short circuit reduction for any
+            # unused parameters. Only if `find_unused_parameters` is set.
+            if self.find_unused_parameters:
+                self.reducer.prepare_for_backward(list(_find_tensors(output)))
+            else:
+                self.reducer.prepare_for_backward([])
+        return output
+
+
 def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
     r"""Applies each `module` in :attr:`modules` in parallel on arguments
     contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
@@ -102,4 +161,4 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
         if isinstance(output, Exception):
             raise output
         outputs.append(output)
-    return outputs
+    return outputs

pytorch_lightning/root_module/hooks.py

@@ -19,3 +19,6 @@ class ModelHooks(torch.nn.Module):
     def on_post_performance_check(self):
         pass
 
+    def on_tng_metrics(self, metrics):
+        pass
+

pytorch_lightning/root_module/model_saving.py

@@ -2,7 +2,7 @@ import torch
 import os
 import re
 import pdb
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 
 class ModelIO(object):
 
@@ -21,6 +21,23 @@ class ModelIO(object):
         """
         raise NotImplementedError
 
+    # -------------------------
+    # OPTIONAL HOOKS
+    # -------------------------
+    def on_hpc_save(self):
+        """
+        Hook to do whatever you need right before Slurm manager saves the model
+        :return:
+        """
+        pass
+
+    def on_hpc_load(self):
+        """
+        Hook to do whatever you need right before Slurm manager loads the model
+        :return:
+        """
+        pass
+
 
 class TrainerIO(object):
 
@@ -49,7 +66,7 @@ class TrainerIO(object):
         checkpoint['optimizer_states'] = optimizer_states
 
         # request what to save from the model
-        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        model = self.model.module if type(self.model) is LightningDistributedDataParallel else self.model
         checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items
@@ -116,6 +133,9 @@ class TrainerIO(object):
             os.makedirs(folderpath, exist_ok=True)
         filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
 
+        # give model a chance to do something on hpc_save
+        self.on_hpc_save()
+
         # request what to save from the model
         checkpoint_dict = self.dump_checkpoint()
 
@@ -137,6 +157,9 @@ class TrainerIO(object):
         model = self.model.module if type(self.model) is LightningDataParallel else self.model
         model.load_model_specific(checkpoint)
 
+        # call model hook
+        self.on_hpc_load()
+
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)
         files = [x for x in files if 'ckpt_' in x]

pytorch_lightning/root_module/root_module.py

@@ -9,7 +9,6 @@ from pytorch_lightning.root_module.optimization import OptimizerConfig
 from pytorch_lightning.root_module.hooks import ModelHooks
 
 
-
 class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
 
     def __init__(self, hparams):
@@ -21,29 +20,17 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.current_epoch = 0
         self.global_step = 0
         self.loaded_optimizer_states_dict = {}
-        self.fast_dev_run = hparams.fast_dev_run
-        self.overfit = hparams.overfit
-        self.gradient_clip = hparams.gradient_clip
         self.trainer = None
-        self.from_lightning = True
         self.experiment = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
-        try:
-            self.on_gpu = hparams.on_gpu
-        except Exception as e:
-            pass
 
         # computed vars for the dataloaders
         self._tng_dataloader = None
         self._val_dataloader = None
         self._test_dataloader = None
 
-        if self.on_gpu:
-            print('running on gpu...')
-            torch.set_default_tensor_type(hparams.default_tensor_type)
-
     def forward(self, *args, **kwargs):
         """
         Expand model in into whatever you need.
@@ -91,7 +78,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         :param logs:
         :return:
         """
-        raise NotImplementedError
+        return logs
 
     def loss(self, *args, **kwargs):
         """
@@ -105,9 +92,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         model_summary = ModelSummary(self)
         print(model_summary)
 
-    def nb_batches(self, dataloader):
-        a = math.ceil(float(len(dataloader.dataset) / self.batch_size))
-        return int(a)
 
     def freeze(self):
         for param in self.parameters():
@@ -145,7 +129,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
     def get_process_position(gpus):
         try:
             current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
-            gpu_ids = gpus.split(';')
+            gpu_ids = gpus.split(',')
             process_position = gpu_ids.index(current_gpu)
             return process_position, current_gpu
         except Exception as e:

pytorch_lightning/utils/arg_parse.py

@@ -39,8 +39,8 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
         parser.add_argument('--model_name', default='', help=','.join(possible_model_names))
 
     # test_tube settings
-    parser.add_argument('-en', '--tt_name', default='r_lib_')
+    parser.add_argument('-en', '--tt_name', default='pt_test')
-    parser.add_argument('-td', '--tt_description', default='test research  lib')
+    parser.add_argument('-td', '--tt_description', default='pytorch lightning test')
     parser.add_argument('--tt_save_path', default=root_dir + '/test_tube_logs', help='logging dir')
     parser.add_argument('--enable_single_run', dest='single_run', action='store_true')
     parser.add_argument('--nb_hopt_trials', default=1, type=int)

pytorch_lightning/utils/plotting.py

@@ -1,4 +1,3 @@
-from matplotlib import pyplot as plt
 import numpy as np
 np.seterr(divide='ignore', invalid='ignore')
 
@@ -13,6 +12,7 @@ def plot_confusion_matrix(cm,
     This function prints and plots the confusion matrix.
     Normalization can be applied by setting `normalize=True`.
     """
+    from matplotlib import pyplot as plt
     if normalize:
         cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
         print("Normalized confusion matrix")

requirements.txt

@@ -1,41 +1,32 @@
-absl-py==0.7.1
+
-astor==0.8.0
+atomicwrites==1.2.1
-bleach==3.1.0
+attrs==18.2.0
-certifi==2019.6.16
+certifi==2018.11.29
-cffi==1.12.3
+cffi==1.11.5
-chardet==3.0.4
+imageio==2.4.1
-docutils==0.14
+mkl-fft==1.0.6
-gast==0.2.2
-google-pasta==0.1.7
-grpcio==1.21.1
-h5py==2.9.0
-idna==2.8
-imageio==2.5.0
-Keras-Applications==1.0.8
-Keras-Preprocessing==1.1.0
-Markdown==3.1.1
-mkl-fft==1.0.12
 mkl-random==1.0.2
-numpy==1.16.4
+more-itertools==5.0.0
+numpy==1.15.4
 olefile==0.46
-pandas==0.24.2
+pandas==0.23.4
-Pillow==6.0.0
+Pillow==5.3.0
-pkginfo==1.5.0.1
+pluggy==0.8.0
-protobuf==3.8.0
+py==1.7.0
 pycparser==2.19
-Pygments==2.4.1
+pytest==4.0.2
-python-dateutil==2.8.0
+python-dateutil==2.7.5
-pytz==2019.1
+pytz==2018.7
-readme-renderer==24.0
+scikit-learn==0.20.2
-requests==2.22.0
+scipy==1.2.0
-requests-toolbelt==0.9.1
 six==1.12.0
+sklearn==0.0
 tensorboard==1.14.0
 tensorboardX==1.7
 tensorflow==1.14.0
-tensorflow-estimator==1.14.0
-termcolor==1.1.0
 test-tube==0.643
+torch==1.0.0
+torchvision==0.2.1
 tqdm==4.32.1
 twine==1.13.0
 urllib3==1.25.3

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.113',
+    version='0.2.4.1',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -17,10 +17,9 @@ setup(
     keywords=["deep learning", "pytorch", "AI"],
     python_requires=">=3.5",
     install_requires=[
-        "torch>=1.0.0",
+        "torch>=1.1.0",
         "tqdm",
-        "test-tube>=0.643",
+        "test-tube>=0.6.7.1",
-        "tensorflow>=1.14.0"
     ],
     packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),