release v0.21

.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)
-    loss = my_loss(out, y)
+    hidden_states = self.encoder(x)
+     
+    # 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
 ```
+   
+## TensorboardX    
+Lightning is fully integrated with tensorboardX.   
 
-## 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
+...   
+```   
+

docs/index.md

@@ -1,5 +1,3 @@
-# PYTORCH-LIGHTNING DOCUMENTATION
-
 ###### New project Quick Start    
 To start a new project define these two files.   
 

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,
-            shuffle=True
+            batch_size=batch_size,
+            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('--out_features', default=10)
-        parser.add_argument('--hidden_dim', default=50000) # use 500 for CPU, 50000 for GPU to see speed difference
+        parser.add_argument('--in_features', default=28*28, type=int)
+        parser.add_argument('--out_features', default=10, type=int)
+        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_cpu_template.py

@@ -41,7 +41,6 @@ def main(hparams):
     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',

examples/new_project_templates/trainer_gpu_cluster_template.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 examples.new_project_templates.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',
-        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/__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
-import numpy as np
-from pytorch_lightning.root_module.memory import get_gpu_memory_map
+"""
+The trainer handles all the logic for running a val loop, training loop, distributing, etc...
+"""
+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 pdb
+import torch.multiprocessing as mp
+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,19 +27,6 @@ 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,
@@ -37,6 +36,7 @@ class Trainer(TrainerIO):
                  cluster=None,
                  process_position=0,
                  current_gpu_name=0,
+                 nb_gpu_nodes=1,
                  gpus=None,
                  progress_bar=True,
                  overfit_pct=0.0,
@@ -54,6 +54,7 @@ class Trainer(TrainerIO):
                  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
@@ -79,8 +80,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 = gpus is not None and len(gpus) > 0
+        self.data_parallel_device_ids = None
+        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) > 0
+
+        # process info
+        self.proc_rank = 0
 
         # training state
         self.optimizers = None
@@ -103,11 +122,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
@@ -128,16 +157,18 @@ class Trainer(TrainerIO):
     def __tng_tqdm_dic(self):
         tqdm_dic = {
             'tng_loss': '{0:.3f}'.format(self.avg_loss),
-            'gpu': '{}'.format(self.current_gpu_name),
             'v_nb': '{}'.format(self.experiment.version),
             'epoch': '{}'.format(self.current_epoch),
             'batch_nb':'{}'.format(self.batch_nb),
         }
         tqdm_dic.update(self.tqdm_metrics)
 
+        if self.on_gpu:
+            tqdm_dic['gpu'] = '{}'.format(self.current_gpu_name)
+
         return tqdm_dic
 
-    def __layout_bookeeping(self, model):
+    def __layout_bookeeping(self):
         # training bookeeping
         self.total_batch_nb = 0
         self.running_loss = []
@@ -146,17 +177,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
@@ -177,12 +208,12 @@ class Trainer(TrainerIO):
         :param max_batches: Scalar
         :return:
         """
-        print('validating...')
+        if self.proc_rank == 0:
+            print('validating...')
 
         # enable eval mode
         model.zero_grad()
         model.eval()
-        model.from_lightning = True
 
         # disable gradients to save memory
         torch.set_grad_enabled(False)
@@ -205,7 +236,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)
 
@@ -239,30 +269,146 @@ class Trainer(TrainerIO):
         self.test_dataloader = model.test_dataloader
         self.val_dataloader = model.val_dataloader
 
+        if self.on_gpu and type(self.tng_dataloader.sampler) is not 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):
-
-        model.trainer = self
-
-        # transfer data loaders from model
-        self.__get_dataloaders(model)
-
-        # init training constants
-        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, optimizer = amp.initialize(
-                model, self.optimizers[0], opt_level=self.amp_level,
+            model, optimizers = amp.initialize(
+                model, self.optimizers, opt_level=self.amp_level,
             )
-            self.optimizers[0] = optimizer
-            model.trainer = self
+            self.optimizers = optimizers
+
+        # when using gpus, first thing we do is spawn a new process between each worker
+        # applies to single gpu, multi-gpu and multi-nodes
+        if self.on_gpu:
+            self.experiment = self.experiment.get_meta_copy()
+            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+        else:
+            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
+        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
+        ip = self.__get_root_node_ip(self.proc_rank, self.nb_gpu_nodes)
+        dist.init_process_group("nccl", init_method=f'tcp://{ip}:12001', rank=self.proc_rank, world_size=self.world_size)
+
+        # copy model to each gpu
+        torch.cuda.set_device(gpu_nb)
+        model.cuda(gpu_nb)
+        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb])
+
+        # continue training routine
+        self.__run_pretrain_routine(model)
+
+    def __get_root_node_ip(self, world_gpu_nb, nb_gpu_nodes):
+        """
+        Resolves the ip address of proc 0.
+        Proc 0 writes address to a file. Every other process waits until the ip is available before it starts
+
+        :param world_gpu_nb: gpu number amongst all the world gpus
+        :param nb_gpu_nodes:
+        :param ip_file_dir:
+        :return:
+        """
+        # on one node we use localhost
+        if nb_gpu_nodes == 1:
+            return '127.0.0.1'
+
+        # where to store ip_table
+        ip_file_dir = os.path.join(self.cluster.log_path, 'ip_tables')
+
+        # the first gpu in the world becomes the host
+        # this is based on its global rank
+        # it communicates its ip by saving an ip_table to the slurm cluster logging dir
+        # every other process waits for this ip to appear before continuing
+        ip_table_name = f'.ip_meta_' + os.environ['SLURM_JOB_ID']
+        ip_file = os.path.join(ip_file_dir, ip_table_name)
+        os.makedirs(ip_file_dir, exist_ok=True)
+
+        if world_gpu_nb == 0:
+            # get the proc 0 IP
+            root_ip = subprocess.run(['hostname', '-I'], stdout=subprocess.PIPE).stdout.decode('utf-8')
+            root_ip = root_ip.split(' ')[0]
+
+            # save the ip to the file
+            with open(file=ip_file, mode='w') as f:
+                f.write(root_ip)
+
+            return root_ip
+        else:
+            # wait up to 120 seconds until proc 0 writes
+            # once written, read proc 0's address and use it to configure server
+            for i in range(0, 120):
+                sleep(1.0)
+                if os.path.exists(ip_file):
+                    ip = list(open(file=ip_file, mode='r'))[0]
+                    return ip
+
+    def __run_pretrain_routine(self, model):
+        """
+        Sanity check a few things before starting actual training
+        :param model:
+        :return:
+        """
+        ref_model = model
+        if self.on_gpu:
+            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:
@@ -271,17 +417,19 @@ class Trainer(TrainerIO):
                 self.lr_schedulers.append(scheduler)
 
         # print model summary
-        model.summarize()
+        if self.proc_rank == 0:
+            ref_model.summarize()
 
-        # put on gpu if needed
-        if self.on_gpu:
-            model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+        # give model convenience properties
+        ref_model.trainer = self
+        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:
@@ -344,7 +492,8 @@ 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:
@@ -369,8 +518,10 @@ class Trainer(TrainerIO):
                         metrics.update(grad_norm_dic)
 
                     # log metrics
-                    self.experiment.log(metrics)
-                    self.experiment.save()
+                    scalar_metrics = self.__metrics_to_scalars(metrics, blacklist=self.__log_vals_blacklist())
+                    if self.proc_rank == 0:
+                        self.experiment.log(scalar_metrics, global_step=self.global_step)
+                        self.experiment.save()
 
                 # hook
                 if self.__is_function_implemented('on_batch_end'):
@@ -396,6 +547,24 @@ class Trainer(TrainerIO):
                 if stop:
                     return
 
+    def __metrics_to_scalars(self, metrics, blacklist=[]):
+        new_metrics = {}
+        for k, v in metrics.items():
+            if type(v) is torch.Tensor:
+                v = v.item()
+
+            if type(v) is dict:
+                v = self.__metrics_to_scalars(v)
+
+            if k not in blacklist:
+                new_metrics[k] = float(v)
+
+        return new_metrics
+
+    def __log_vals_blacklist(self):
+        """avoid logging some vals lightning uses to maintain state"""
+        blacklist = {'batch_nb', 'v_nb', 'epoch', 'gpu'}
+        return blacklist
 
     def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
@@ -416,7 +585,6 @@ 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)
 
@@ -427,6 +595,7 @@ class Trainer(TrainerIO):
 
         # 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()
@@ -512,5 +681,6 @@ class Trainer(TrainerIO):
             self.prog_bar.set_postfix(**tqdm_metrics)
 
         # model checkpointing
-        print('save callback...')
-        self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
+        if self.proc_rank == 0:
+            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/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
@@ -107,12 +124,18 @@ class TrainerIO(object):
         # save exp to make sure we get all the metrics
         experiment.save()
 
+        # close experiment to avoid issues
+        experiment.close()
+
         ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
 
         if not os.path.exists(folderpath):
             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()
 
@@ -134,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,28 +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.
@@ -104,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():

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('-td', '--tt_description', default='test research  lib')
+    parser.add_argument('-en', '--tt_name', default='pt_test')
+    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)

requirements.txt

@@ -1,8 +1,8 @@
+
 atomicwrites==1.2.1
 attrs==18.2.0
 certifi==2018.11.29
 cffi==1.11.5
-h5py==2.9.0
 imageio==2.4.1
 mkl-fft==1.0.6
 mkl-random==1.0.2
@@ -21,7 +21,15 @@ scikit-learn==0.20.2
 scipy==1.2.0
 six==1.12.0
 sklearn==0.0
-test-tube==0.6282
+tensorboard==1.14.0
+tensorboardX==1.7
+tensorflow==1.14.0
+test-tube==0.643
 torch==1.0.0
 torchvision==0.2.1
-tqdm==4.28.1
+tqdm==4.32.1
+twine==1.13.0
+urllib3==1.25.3
+webencodings==0.5.1
+Werkzeug==0.15.4
+wrapt==1.11.2

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.112',
+    version='0.21',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -17,9 +17,10 @@ setup(
     keywords=["deep learning", "pytorch", "AI"],
     python_requires=">=3.5",
     install_requires=[
-        "torch>=1.0.0",
+        "torch>=1.1.0",
         "tqdm",
-        "test-tube",
+        "test-tube>=0.653",
+        "tensorflow>=1.14.0"
     ],
     packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),