release v0.11

.gitignore

@@ -116,3 +116,6 @@ ENV/
 
 # mypy
 .mypy_cache/
+
+# data
+mnist/

README.md

@@ -55,6 +55,29 @@ To use lightning do 2 things:
 1. [Define a trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/basic_trainer.py) (which will run ALL your models).   
 2. [Define a model](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/example_model.py).     
 
+#### Quick demo
+Run the following demo to see how it works:
+```bash
+# install lightning
+pip install pytorch-lightning
+
+# clone lightning for the demo
+git clone https://github.com/williamFalcon/pytorch-lightning.git
+cd pytorch-lightning/docs/source/examples
+
+# run demo (on cpu)
+python fully_featured_trainer.py
+```
+
+Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.
+```bash
+# run a grid search on two gpus
+python fully_featured_trainer.py --gpus "0;1"
+
+# run single model on multiple gpus
+python fully_featured_trainer.py --gpus "0;1" --interactive
+```
+
 #### Basic trainer example  
 See [this demo](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/fully_featured_trainer.py) for a more robust trainer example.   
 

docs/source/examples/__init__.py

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

docs/source/examples/basic_trainer.py

@@ -2,10 +2,10 @@ import os
 import sys
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch-lightning.models.trainer import Trainer
-from pytorch-lightning.utils.arg_parse import add_default_args
-from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-from demo.example_model import ExampleModel
+from pytorch_lightning.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
+from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+from docs.source.examples.example_model import ExampleModel
 
 
 def main(hparams):

docs/source/examples/example_model.py

@@ -1,11 +1,13 @@
 import torch.nn as nn
 import numpy as np
-from pytorch-lightning.root_module.root_module import RootModule
+from pytorch_lightning.root_module.root_module import RootModule
 from test_tube import HyperOptArgumentParser
 from torchvision.datasets import MNIST
 import torchvision.transforms as transforms
 import torch
 import torch.nn.functional as F
+import os, pdb
+from collections import OrderedDict
 
 
 class ExampleModel(RootModule):
@@ -40,8 +42,9 @@ class ExampleModel(RootModule):
     # TRAINING
     # ---------------------
     def forward(self, x):
+
         x = self.c_d1(x)
-        x = F.tanh(x)
+        x = torch.tanh(x)
         x = self.c_d1_bn(x)
         x = self.c_d1_drop(x)
 
@@ -54,7 +57,7 @@ class ExampleModel(RootModule):
         nll = F.nll_loss(logits, labels)
         return nll
 
-    def training_step(self, data_batch):
+    def training_step(self, data_batch, batch_i):
         """
         Called inside the training loop
         :param data_batch:
@@ -68,10 +71,16 @@ class ExampleModel(RootModule):
         # calculate loss
         loss_val = self.loss(y, y_hat)
 
-        tqdm_dic = {'tng_loss': loss_val.item()}
-        return loss_val, tqdm_dic
+        # tqdm_dic = {'tng_loss': loss_val.item()}
+        # return loss_val, tqdm_dic
 
-    def validation_step(self, data_batch):
+        output = OrderedDict({
+            'loss': loss_val,
+            'tqdm_metrics': {}
+        })
+        return output
+
+    def validation_step(self, data_batch, batch_i):
         """
         Called inside the validation loop
         :param data_batch:
@@ -87,9 +96,14 @@ class ExampleModel(RootModule):
         labels_hat = torch.argmax(y_hat, dim=1)
         val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
 
-        output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
+        # output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
+        output = OrderedDict({
+            'val_loss': loss_val,
+            'val_acc': torch.tensor(val_acc),
+        })
         return output
 
+
     def validation_end(self, outputs):
         """
         Called at the end of validation to aggregate outputs
@@ -97,13 +111,14 @@ class ExampleModel(RootModule):
         :return:
         """
         val_loss_mean = 0
-        accs = []
+        val_acc_mean = 0
         for output in outputs:
             val_loss_mean += output['val_loss']
-            accs.append(output['val_acc'])
+            val_acc_mean += output['val_acc']
 
         val_loss_mean /= len(outputs)
-        tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)}
+        val_acc_mean /= len(outputs)
+        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
 
     def update_tng_log_metrics(self, logs):
@@ -177,7 +192,7 @@ class ExampleModel(RootModule):
         return self._test_dataloader
 
     @staticmethod
-    def add_model_specific_args(parent_parser):
+    def add_model_specific_args(parent_parser, root_dir):
         parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
 
         # param overwrites
@@ -186,11 +201,11 @@ class ExampleModel(RootModule):
         # 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('--hidden_dim', default=500)
         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
 
         # data
-        parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str)
+        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],

docs/source/examples/fully_featured_trainer.py

@@ -5,10 +5,10 @@ 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.models.trainer import Trainer
+from pytorch_lightning.utils.arg_parse import add_default_args
 
-from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
 
 SEED = 2334
 torch.manual_seed(SEED)
@@ -17,7 +17,7 @@ np.random.seed(SEED)
 # ---------------------
 # DEFINE MODEL HERE
 # ---------------------
-from example_model import ExampleModel
+from docs.source.examples.example_model import ExampleModel
 # ---------------------
 
 AVAILABLE_MODELS = {
@@ -42,9 +42,7 @@ def main(hparams, cluster, results_dict):
     :param hparams:
     :return:
     """
-    on_gpu = torch.cuda.is_available()
-    if hparams.disable_cuda:
-        on_gpu = False
+    on_gpu = hparams.gpus is not None and torch.cuda.is_available()
 
     device = 'cuda' if on_gpu else 'cpu'
     hparams.__setattr__('device', device)
@@ -93,12 +91,18 @@ def main(hparams, cluster, results_dict):
         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
@@ -159,35 +163,40 @@ if __name__ == '__main__':
         model_name = 'model_template'
 
     # use default args
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
+    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)
+    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"
-    gpu_ids = hyperparams.gpus.split(';')
 
+    # ---------------------
     # RUN TRAINING
+    # ---------------------
+
+    # cluster and CPU
     if hyperparams.on_cluster:
-        # Gets called when running via HPC 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.single_run_gpu:
-        # run on 1 gpu
-        print(f'RUNNING 1 TRIAL ON GPU. gpu: {gpu_ids[0]}')
-        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[0]
+    elif hyperparams.gpus is None:
+        # run on cpu
+        print('RUNNING ON CPU')
         main(hyperparams, None, None)
 
-    elif hyperparams.local or hyperparams.single_run:
-        # run 1 trial but on CPU
-        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
-        print('RUNNING LOCALLY')
+    # 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:

pytorch_lightning/models/trainer.py

@@ -5,6 +5,27 @@ from pytorch_lightning.root_module.memory import get_gpu_memory_map
 import traceback
 from pytorch_lightning.root_module.model_saving import TrainerIO
 from torch.optim.lr_scheduler import MultiStepLR
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
+import pdb
+
+try:
+    from apex import amp
+    APEX_AVAILABLE = True
+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):
@@ -15,7 +36,7 @@ class Trainer(TrainerIO):
                  cluster=None,
                  process_position=0,
                  current_gpu_name=0,
-                 on_gpu=False,
+                 gpus=None,
                  enable_tqdm=True,
                  overfit_pct=0.0,
                  track_grad_norm=-1,
@@ -26,6 +47,9 @@ class Trainer(TrainerIO):
                  train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0, val_check_interval=0.95,
                  log_save_interval=1, add_log_row_interval=1,
                  lr_scheduler_milestones=None,
+                 use_amp=False,
+                 check_grad_nans=False,
+                 amp_level='O2',
                  nb_sanity_val_steps=5):
 
         # Transfer params
@@ -33,7 +57,7 @@ class Trainer(TrainerIO):
         self.enable_early_stop = enable_early_stop
         self.track_grad_norm = track_grad_norm
         self.fast_dev_run = fast_dev_run
-        self.on_gpu = on_gpu
+        self.on_gpu = gpus is not None and torch.cuda.is_available()
         self.enable_tqdm = enable_tqdm
         self.experiment = experiment
         self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
@@ -51,6 +75,10 @@ class Trainer(TrainerIO):
         self.nb_sanity_val_steps = nb_sanity_val_steps
         self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')]
         self.lr_schedulers = []
+        self.amp_level = amp_level
+        self.check_grad_nans = check_grad_nans
+        self.data_parallel_device_ids = gpus
+        self.data_parallel = gpus is not None and len(gpus) > 0
 
         # training state
         self.optimizers = None
@@ -73,6 +101,11 @@ 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
+        self.use_amp = use_amp and APEX_AVAILABLE
+        if self.use_amp:
+            print('using 16bit precision')
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -101,7 +134,7 @@ class Trainer(TrainerIO):
         tqdm_dic.update(self.tqdm_metrics)
         return tqdm_dic
 
-    def __layout_bookeeping(self):
+    def __layout_bookeeping(self, model):
         # training bookeeping
         self.total_batch_nb = 0
         self.running_loss = []
@@ -110,21 +143,21 @@ class Trainer(TrainerIO):
         self.tqdm_metrics = {}
 
         # determine number of training batches
-        nb_tng_batches = self.model.nb_batches(self.tng_dataloader)
-        self.nb_tng_batches = int(nb_tng_batches * self.train_percent_check)
+        self.nb_tng_batches = model.nb_batches(self.tng_dataloader)
+        self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
 
         # determine number of validation batches
-        nb_val_batches = self.model.nb_batches(self.val_dataloader)
-        nb_val_batches = int(nb_val_batches * self.val_percent_check)
-        nb_val_batches = max(1, nb_val_batches)
-        self.nb_val_batches = nb_val_batches
+        self.nb_val_batches = model.nb_batches(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
-        nb_test_batches = self.model.nb_batches(self.test_dataloader)
-        self.nb_test_batches = int(nb_test_batches * self.test_percent_check)
+        self.nb_test_batches = model.nb_batches(self.test_dataloader)
+        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
 
         # determine when to check validation
-        self.val_check_batch = int(nb_tng_batches * self.val_check_interval)
+        self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval)
 
     def __add_tqdm_metrics(self, metrics):
         for k, v in metrics.items():
@@ -143,6 +176,7 @@ class Trainer(TrainerIO):
         # enable eval mode
         model.zero_grad()
         model.eval()
+        model.from_lightning = True
 
         # disable gradients to save memory
         torch.set_grad_enabled(False)
@@ -151,19 +185,24 @@ class Trainer(TrainerIO):
         outputs = []
 
         # run training
-        for i, data_batch in enumerate(dataloader):
+        for batch_i, data_batch in enumerate(dataloader):
 
             if data_batch is None:
                 continue
 
             # stop short when on fast dev run
-            if max_batches is not None and i >= max_batches:
+            if max_batches is not None and batch_i >= max_batches:
                 break
 
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
-            output = model.validation_step(data_batch)
+            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)
+
             outputs.append(output)
 
             # batch done
@@ -171,13 +210,17 @@ class Trainer(TrainerIO):
                 self.prog_bar.update(1)
 
         # give model a chance to do something with the outputs
-        val_results = model.validation_end(outputs)
+        if self.data_parallel:
+            val_results = model.module.validation_end(outputs)
+        else:
+            val_results = model.validation_end(outputs)
 
         # enable train mode again
         model.train()
 
         # enable gradients to save memory
         torch.set_grad_enabled(True)
+
         return val_results
 
     def __get_dataloaders(self, model):
@@ -194,19 +237,27 @@ class Trainer(TrainerIO):
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
-        self.model = model
+
         model.trainer = self
 
         # transfer data loaders from model
         self.__get_dataloaders(model)
 
         # init training constants
-        self.__layout_bookeeping()
+        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()
 
+        if self.use_amp:
+            # An example
+            model, optimizer = amp.initialize(
+                model, self.optimizers[0], opt_level=self.amp_level,
+            )
+            self.optimizers[0] = optimizer
+            model.trainer = self
+
         # add lr schedulers
         if self.lr_scheduler_milestones is not None:
             for optimizer in self.optimizers:
@@ -218,7 +269,7 @@ class Trainer(TrainerIO):
 
         # put on gpu if needed
         if self.on_gpu:
-            model = model.cuda()
+            model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
 
         # 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)
@@ -233,6 +284,7 @@ class Trainer(TrainerIO):
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
+        self.model = model
         self.__train()
 
     def __train(self):
@@ -242,11 +294,13 @@ class Trainer(TrainerIO):
             for lr_scheduler in self.lr_schedulers:
                 lr_scheduler.step()
 
-            self.model.current_epoch = epoch_nb
+            model = self.model.module if self.data_parallel else self.model
+            model.current_epoch = epoch_nb
 
             # hook
             if self.__is_function_implemented('on_epoch_start'):
-                self.model.on_epoch_start()
+                model = self.model.module if self.data_parallel else self.model
+                model.on_epoch_start()
 
             self.current_epoch = epoch_nb
             self.total_batches = self.nb_tng_batches + self.nb_val_batches
@@ -259,7 +313,9 @@ class Trainer(TrainerIO):
             for batch_nb, data_batch in enumerate(self.tng_dataloader):
                 self.batch_nb = batch_nb
                 self.global_step += 1
-                self.model.global_step = self.global_step
+
+                model = self.model.module if self.data_parallel else self.model
+                model.global_step = self.global_step
 
                 # stop when the flag is changed or we've gone past the amount requested in the batches
                 self.total_batch_nb += 1
@@ -270,7 +326,7 @@ class Trainer(TrainerIO):
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
-                batch_result = self.__run_tng_batch(data_batch)
+                batch_result = self.__run_tng_batch(data_batch, batch_nb)
                 early_stop_epoch = batch_result == -1
 
                 # ---------------
@@ -281,15 +337,18 @@ class Trainer(TrainerIO):
                     self.__run_validation()
 
                 # when batch should be saved
-                if (batch_nb + 1) % self.log_save_interval == 0:
+                if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
                     self.experiment.save()
 
                 # when metrics should be logged
-                if batch_nb % self.add_log_row_interval == 0:
+                if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
-                    metrics = self.model.update_tng_log_metrics(self.__tng_tqdm_dic)
+                    if self.data_parallel:
+                        metrics = self.model.module.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:
@@ -298,7 +357,9 @@ class Trainer(TrainerIO):
 
                     # add norms
                     if self.track_grad_norm > 0:
-                        grad_norm_dic = self.model.grad_norm(self.track_grad_norm)
+                        model = self.model.module if self.data_parallel else self.model
+                        grad_norm_dic = model.grad_norm(self.track_grad_norm)
+
                         metrics.update(grad_norm_dic)
 
                     # log metrics
@@ -307,7 +368,8 @@ class Trainer(TrainerIO):
 
                 # hook
                 if self.__is_function_implemented('on_batch_end'):
-                    self.model.on_batch_end()
+                    model = self.model.module if self.data_parallel else self.model
+                    model.on_batch_end()
 
                 # end epoch early
                 if early_stop_epoch:
@@ -315,7 +377,8 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_epoch_end'):
-                self.model.on_epoch_end()
+                model = self.model.module if self.data_parallel else self.model
+                model.on_epoch_end()
 
             # early stopping
             if self.enable_early_stop:
@@ -328,13 +391,15 @@ class Trainer(TrainerIO):
                     return
 
 
-    def __run_tng_batch(self, data_batch):
+    def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
             return 0
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            response = self.model.on_batch_start(data_batch)
+            model = self.model.module if self.data_parallel else self.model
+            response = model.on_batch_start(data_batch)
+
             if response == -1:
                 return -1
 
@@ -343,11 +408,30 @@ class Trainer(TrainerIO):
 
         # forward pass
         # return a scalar value and a dic with tqdm metrics
-        loss, model_specific_tqdm_metrics_dic = self.model.training_step(data_batch)
+        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']
+        loss = output['loss']
+
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
-        loss.backward()
+        if self.use_amp:
+            for optimizer in self.optimizers:
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+        else:
+            loss.backward()
+
+        if self.check_grad_nans:
+            model = self.model.module if self.data_parallel else self.model
+            for param in model.parameters():
+                print(param.grad.float().sum())
+
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -0,0 +1,105 @@
+from torch.nn import DataParallel
+
+import threading
+import torch
+from torch.cuda._utils import _get_device_index
+import pdb
+
+
+def get_a_var(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj
+
+    if isinstance(obj, list) or isinstance(obj, tuple):
+        for result in map(get_a_var, obj):
+            if isinstance(result, torch.Tensor):
+                return result
+    if isinstance(obj, dict):
+        for result in map(get_a_var, obj.items()):
+            if isinstance(result, torch.Tensor):
+                return result
+    return None
+
+
+class LightningDataParallel(DataParallel):
+    """
+    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 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)
+    on each of :attr:`devices`.
+
+    Args:
+        modules (Module): modules to be parallelized
+        inputs (tensor): inputs to the modules
+        devices (list of int or torch.device): CUDA devices
+
+    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
+    :attr:`devices` (if given) should all have same length. Moreover, each
+    element of :attr:`inputs` can either be a single object as the only argument
+    to a module, or a collection of positional arguments.
+    """
+    assert len(modules) == len(inputs)
+    if kwargs_tup is not None:
+        assert len(modules) == len(kwargs_tup)
+    else:
+        kwargs_tup = ({},) * len(modules)
+    if devices is not None:
+        assert len(modules) == len(devices)
+    else:
+        devices = [None] * len(modules)
+    devices = list(map(lambda x: _get_device_index(x, True), devices))
+    lock = threading.Lock()
+    results = {}
+    grad_enabled = torch.is_grad_enabled()
+
+    def _worker(i, module, input, kwargs, device=None):
+        torch.set_grad_enabled(grad_enabled)
+        if device is None:
+            device = get_a_var(input).get_device()
+        try:
+            with torch.cuda.device(device):
+                # this also avoids accidental slicing of `input` if it is a Tensor
+                if not isinstance(input, (list, tuple)):
+                    input = (input,)
+
+                # ---------------
+                # CHANGE
+                if module.training:
+                    output = module.training_step(*input, **kwargs)
+                else:
+                    output = module.validation_step(*input, **kwargs)
+                # ---------------
+
+            with lock:
+                results[i] = output
+        except Exception as e:
+            with lock:
+                results[i] = e
+
+    if len(modules) > 1:
+        threads = [threading.Thread(target=_worker,
+                                    args=(i, module, input, kwargs, device))
+                   for i, (module, input, kwargs, device) in
+                   enumerate(zip(modules, inputs, kwargs_tup, devices))]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+    else:
+        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])
+
+    outputs = []
+    for i in range(len(inputs)):
+        output = results[i]
+        if isinstance(output, Exception):
+            raise output
+        outputs.append(output)
+    return outputs

pytorch_lightning/root_module/model_saving.py

@@ -1,7 +1,8 @@
 import torch
 import os
 import re
-
+import pdb
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
 
 class ModelIO(object):
 
@@ -48,7 +49,8 @@ class TrainerIO(object):
         checkpoint['optimizer_states'] = optimizer_states
 
         # request what to save from the model
-        checkpoint_dict = self.model.get_save_dict()
+        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items
         checkpoint.update(checkpoint_dict)
@@ -88,6 +90,7 @@ class TrainerIO(object):
         self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
         self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
         self.global_step = checkpoint['global_step']
+        self.current_epoch = checkpoint['epoch']
 
         # restore the optimizers
         optimizer_states = checkpoint['optimizer_states']
@@ -98,6 +101,9 @@ class TrainerIO(object):
     # PRIVATE OPS
     # ----------------------------------
     def hpc_save(self, folderpath, experiment):
+        # make sure the checkpoint folder exists
+        os.makedirs(folderpath, exist_ok=True)
+
         # save exp to make sure we get all the metrics
         experiment.save()
 
@@ -125,10 +131,15 @@ class TrainerIO(object):
         self.restore_training_state(checkpoint)
 
         # load model state
-        self.model.load_model_specific(checkpoint)
+        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        model.load_model_specific(checkpoint)
 
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)
+        files = [x for x in files if 'ckpt_' in x]
+        if len(files) == 0:
+            return 0
+
         ckpt_vs = []
         for name in files:
             name = name.split('ckpt_')[-1]

pytorch_lightning/root_module/root_module.py

@@ -25,6 +25,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.gradient_clip = hparams.gradient_clip
         self.num = 2
         self.trainer = None
+        self.from_lightning = True
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
@@ -40,8 +41,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
 
         if self.on_gpu:
             print('running on gpu...')
-            self.dtype = torch.cuda.FloatTensor
-            torch.set_default_tensor_type('torch.cuda.FloatTensor')
+            torch.set_default_tensor_type(hparams.default_tensor_type)
 
     def forward(self, *args, **kwargs):
         """
@@ -52,7 +52,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def validation_step(self, data_batch):
+    def validation_step(self, data_batch, batch_nb):
         """
         return whatever outputs will need to be aggregated in validation_end
         :param data_batch:
@@ -68,7 +68,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def training_step(self, data_batch):
+    def training_step(self, data_batch, batch_nb):
         """
         return loss, dict with metrics for tqdm
         :param data_batch:
@@ -151,19 +151,23 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
             return 0, 0
 
     @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu):
+    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
         """
         Primary way of loading model from csv weights path
         :param weights_path:
         :param tags_csv:
         :param on_gpu:
+        :param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
         :return:
         """
         hparams = load_hparams_from_tags_csv(tags_csv)
         hparams.__setattr__('on_gpu', on_gpu)
 
         if on_gpu:
-            checkpoint = torch.load(weights_path)
+            if map_location is not None:
+                checkpoint = torch.load(weights_path, map_location=map_location)
+            else:
+                checkpoint = torch.load(weights_path)
         else:
             checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 

pytorch_lightning/utils/arg_parse.py

@@ -1,3 +1,5 @@
+import pdb
+
 def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None):
 
     # tng, test, val check intervals
@@ -45,10 +47,13 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     parser.add_argument('--log_stdout', dest='log_stdout', action='store_true')
 
     # GPU
-    parser.add_argument('--per_experiment_nb_gpus', default=1, type=int)
-    parser.add_argument('--gpus', default='0', type=str)
+    parser.add_argument('--gpus', default=None, type=str)
     parser.add_argument('--single_run_gpu', dest='single_run_gpu', action='store_true')
-    parser.add_argument('--disable_cuda', dest='disable_cuda', action='store_true')
+    parser.add_argument('--default_tensor_type', default='torch.cuda.FloatTensor', type=str)
+    parser.add_argument('--use_amp', dest='use_amp', action='store_true')
+    parser.add_argument('--check_grad_nans', dest='check_grad_nans', action='store_true')
+    parser.add_argument('--amp_level', default='O2',type=str)
+
 
     # run on hpc
     parser.add_argument('--on_cluster', dest='on_cluster', action='store_true')
@@ -63,9 +68,9 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     if rand_seed is not None:
         parser.add_argument('--random_seed', default=rand_seed, type=int)
 
-    parser.add_argument('--live', dest='live', action='store_true', help='runs on gpu without cluster')
-    parser.add_argument('--enable_debug', dest='debug', action='store_true', help='enables/disables test tube')
-    parser.add_argument('--enable_local', dest='local', action='store_true', help='enables local tng')
+    parser.add_argument('--interactive', dest='interactive', action='store_true', help='runs on gpu without cluster')
+    parser.add_argument('--debug', dest='debug', action='store_true', help='enables/disables test tube')
+    parser.add_argument('--local', dest='local', action='store_true', help='enables local tng')
 
     # optimizer
     parser.add_argument('--lr_scheduler_milestones', default=None, type=str)

pytorch_lightning/utils/pt_callbacks.py

@@ -1,5 +1,6 @@
 import numpy as np
 import os, shutil
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDataParallel
 
 
 class Callback(object):
@@ -33,6 +34,8 @@ class Callback(object):
         self.params = params
 
     def set_model(self, model):
+        if type(model) is LightningDataParallel:
+            model = model.module
         self.model = model
 
     def on_epoch_begin(self, epoch, logs=None):

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.1.dev15',
+    version='0.11',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -17,12 +17,13 @@ setup(
     keywords=["deep learning", "pytorch", "AI"],
     python_requires=">=3.5",
     install_requires=[
-        "torch",
+        "torch>=1.0.0",
         "tqdm",
         "test-tube",
     ],
     packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),
+    long_description_content_type='text/markdown',
     include_package_data=True,
     zip_safe=False,
 )