release v0.11

Fix some link bugs in the README.md

.gitignore

@@ -0,0 +1,121 @@
+# project
+.DS_Store
+.data/
+run_configs/
+test_tube_logs/
+test_tube_data/
+datasets/
+model_weights/
+app/models/
+pip-wheel-metadata/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+example.py
+timit_data/
+LJSpeech-1.1/
+
+# C extensions
+*.so
+
+.idea/
+
+# Distribution / packaging
+.Python
+env/
+ide_layouts/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# data
+mnist/

MANIFEST.in

@@ -0,0 +1,9 @@
+graft docs
+
+include COPYING
+include AUTHORS
+
+recursive-include src/einsteinpy/tests *.py *.html
+
+prune docs/source/examples/.ipynb_checkpoints
+global-exclude *.py[cod] __pycache__ *.so *.dylib

README.md

@@ -1,66 +1,150 @@
-# Pytorch-lightning
-The Keras for ML-researchers in PyTorch.    
+<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/lightning_logo.png" width="50">
+  </a>
+</p>
+<h3 align="center">
+  Pytorch Lightning
+</h3>
+<p align="center">
+  The Keras for ML researchers using PyTorch. More control. Less boilerplate.    
+</p>
+<p align="center">
+  <a href="https://badge.fury.io/py/pytorch-lightning"><img src="https://badge.fury.io/py/pytorch-lightning.svg" alt="PyPI version" height="18"></a>
+<!--   <a href="https://travis-ci.org/williamFalcon/test-tube"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a> -->
+  <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/COPYING"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
+</p>   
+
+```bash
+pip install pytorch-lightning    
+```
+
+## Docs   
+In progress. Documenting now!  
+
+## Disclaimer
+This is a research tool I built for myself internally while doing my PhD. The API is not 100% production quality, but my hope is that by open-sourcing, we can all get it there (I don't have too much time nowadays to write production-level code).
+
+## What is it?  
+Keras is too abstract for researchers. Lightning makes it so you only have to define your model but still control all details of training if you need to.   
+
+Pytorch   
+<-- Lightning   
+Your model.   
+
+**Lightning will do the following for you:**   
+
+1. Run the training loop.   
+2. Run the validation loop.   
+3. Run the testing loop.   
+4. Early stopping.   
+5. Learning rate annealing.
+6. Can train complex models like GANs or anything with multiple optimizers.
+7. Weight checkpointing.   
+8. Model saving.   
+9. Model loading.   
+10. Log training details (through test-tube).  
+11. Run training on multiple GPUs (through test-tube).   
+12. Run training on a GPU cluster managed by SLURM (through test-tube).   
+13. Distribute memory-bound models on multiple GPUs.
+14. Give your model hyperparameters parsed from the command line OR a JSON file.   
+15. Run your model in a dev environment where nothing logs.      
 
 ## Usage
 To use lightning do 2 things:  
-1. [Define a trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/trainer_main.py) (which will run ALL your models).   
-2. [Define a model](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/sample_model_template/model_template.py).     
+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).     
 
-### Example:   
+#### Quick demo
+Run the following demo to see how it works:
+```bash
+# install lightning
+pip install pytorch-lightning
 
-#### Define the trainer   
+# clone lightning for the demo
+git clone https://github.com/williamFalcon/pytorch-lightning.git
+cd pytorch-lightning/docs/source/examples
 
-```python
-# trainer.py
-
-from pytorch_lightning.models.trainer import Trainer   
-from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-from my_project import My_Model   
-from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-
-# --------------
-# TEST TUBE INIT
-exp = Experiment(
-    name='my_exp',
-    debug=True,
-    save_dir='/some/path',
-    autosave=False,
-    description='my desc'
-)
-
-# --------------------
-# CALLBACKS
-early_stop = EarlyStopping(
-    monitor='val_loss',
-    patience=3,
-    verbose=True,
-    mode='min'
-)
-
-model_save_path = 'PATH/TO/SAVE'
-checkpoint = ModelCheckpoint(
-    filepath=model_save_path,
-    save_function=None,
-    save_best_only=True,
-    verbose=True,
-    monitor='val_acc',
-    mode='min'
-)
-
-# configure trainer
-trainer = Trainer(
-    experiment=experiment,
-    cluster=cluster,
-    checkpoint_callback=checkpoint,
-    early_stop_callback=early_stop
-)  
-
-# init model and train
-model = My_Model()
-trainer.fit(model)
+# run demo (on cpu)
+python fully_featured_trainer.py
 ```
 
-#### Define the model   
+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.   
+
+```python
+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
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # init experiment
+    exp = Experiment(
+        name=hparams.tt_name,
+        debug=hparams.debug,
+        save_dir=hparams.tt_save_path,
+        version=hparams.hpc_exp_number,
+        autosave=False,
+        description=hparams.tt_description
+    )
+
+    exp.argparse(hparams)
+    exp.save()
+
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+
+    # build model
+    model = ExampleModel(hparams)
+
+    # callbacks
+    early_stop = EarlyStopping(monitor='val_acc', patience=3, mode='min', verbose=True)
+    checkpoint = ModelCheckpoint(filepath=model_save_path, save_function=None, save_best_only=True, verbose=True, monitor='val_acc', mode='min')
+
+    # configure trainer
+    trainer = Trainer(experiment=exp, checkpoint_callback=checkpoint, early_stop_callback=early_stop)
+
+    # train model
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    # use default args given by lightning
+    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
+    parent_parser = HyperOptArgumentParser(strategy='random_search', add_help=False)
+    add_default_args(parent_parser, root_dir)
+
+    # allow model to overwrite or extend args
+    parser = ExampleModel.add_model_specific_args(parent_parser)
+    hyperparams = parser.parse_args()
+
+    # train model
+    main(hyperparams)
+
+```
+
+#### Basic model example  
+Here we only show the method signatures. It's up to you to define the content.   
 
 ```python
 from torch import nn
@@ -69,32 +153,32 @@ class My_Model(RootModule):
     def __init__(self):
         # define model
         self.l1 = nn.Linear(200, 10)
-    
+
     # ---------------
     # TRAINING
     def training_step(self, data_batch):
         x, y = data_batch
         y_hat = self.l1(x)
         loss = some_loss(y_hat)
-        
+
         return loss_val, {'train_loss': loss}
-    
+
     def validation_step(self, data_batch):
         x, y = data_batch
         y_hat = self.l1(x)
         loss = some_loss(y_hat)
-        
+
         return loss_val, {'val_loss': loss}
- 
+
      def validation_end(self, outputs):
         total_accs = []
-        
+
         for output in outputs:
             total_accs.append(output['val_acc'].item())
-        
+
         # return a dict
         return {'total_acc': np.mean(total_accs)}
-     
+
      # ---------------
      # SAVING
      def get_save_dict(self):
@@ -106,7 +190,7 @@ class My_Model(RootModule):
     def load_model_specific(self, checkpoint):
         # lightning loads for you. Here's your chance to say what you want to load
         self.load_state_dict(checkpoint['state_dict'])
-    
+
     # ---------------
     # TRAINING CONFIG
     def configure_optimizers(self):
@@ -114,7 +198,7 @@ class My_Model(RootModule):
         # lightning will call automatically
         optimizer = self.choose_optimizer('adam', self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
         return [optimizer]
-    
+
     @property
     def tng_dataloader(self):
         return pytorch_dataloader('train')
@@ -126,7 +210,7 @@ class My_Model(RootModule):
     @property
     def test_dataloader(self):
         return pytorch_dataloader('test')
-        
+
     # ---------------
     # MODIFY YOUR COMMAND LINE ARGS
     @staticmethod
@@ -135,6 +219,8 @@ class My_Model(RootModule):
         parser.add_argument('--out_features', default=20)
         return parser
 ```
+
+
 ### Details    
 
 #### Model definition   
@@ -143,7 +229,7 @@ class My_Model(RootModule):
 | training_step  |  Called with a batch of data during training | data from your dataloaders  | tuple: scalar, dict  |
 | validation_step  |  Called with a batch of data during validation  | data from your dataloaders  | tuple: scalar, dict |
 | validation_end  |  Collate metrics from all validation steps |  outputs: array where each item is the output of a validation step | dict: for logging |
-| get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved | 
+| get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved |
 
 #### Model training   
 | Name  | Description  |  Input |  Return |
@@ -159,7 +245,7 @@ class My_Model(RootModule):
 |---|---|---|---|
 | get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved |
 |  load_model_specific |  called when loading a model | checkpoint: dict you created in get_save_dict  | dict: modified in whatever way you want  |
-    
+
 ## Optional model hooks.   
 Add these to the model whenever you want to configure training behavior.    
 

docs/source/examples/__init__.py

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

docs/source/examples/basic_trainer.py

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

docs/source/examples/example_model.py

@@ -0,0 +1,215 @@
+import torch.nn as nn
+import numpy as np
+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):
+    """
+    Sample model to show how to define a template
+    """
+
+    def __init__(self, hparams):
+        # init superclass
+        super(ExampleModel, self).__init__(hparams)
+
+        self.batch_size = hparams.batch_size
+
+        # build model
+        self.__build_model()
+
+    # ---------------------
+    # MODEL SETUP
+    # ---------------------
+    def __build_model(self):
+        """
+        Layout model
+        :return:
+        """
+        self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
+        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
+        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
+
+        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features)
+
+    # ---------------------
+    # TRAINING
+    # ---------------------
+    def forward(self, x):
+
+        x = self.c_d1(x)
+        x = torch.tanh(x)
+        x = self.c_d1_bn(x)
+        x = self.c_d1_drop(x)
+
+        x = self.c_d2(x)
+        logits = F.log_softmax(x, dim=1)
+
+        return logits
+
+    def loss(self, labels, logits):
+        nll = F.nll_loss(logits, labels)
+        return nll
+
+    def training_step(self, data_batch, batch_i):
+        """
+        Called inside the training loop
+        :param data_batch:
+        :return:
+        """
+        # forward pass
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        # calculate loss
+        loss_val = self.loss(y, y_hat)
+
+        # tqdm_dic = {'tng_loss': loss_val.item()}
+        # return loss_val, tqdm_dic
+
+        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:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_val = self.loss(y, y_hat)
+
+        # acc
+        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 = 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
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        val_loss_mean = 0
+        val_acc_mean = 0
+        for output in outputs:
+            val_loss_mean += output['val_loss']
+            val_acc_mean += output['val_acc']
+
+        val_loss_mean /= len(outputs)
+        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):
+        return logs
+
+    # ---------------------
+    # MODEL SAVING
+    # ---------------------
+    def get_save_dict(self):
+        checkpoint = {'state_dict': self.state_dict()}
+        return checkpoint
+
+    def load_model_specific(self, checkpoint):
+        self.load_state_dict(checkpoint['state_dict'])
+        pass
+
+    # ---------------------
+    # TRAINING SETUP
+    # ---------------------
+    def configure_optimizers(self):
+        """
+        return whatever optimizers we want here
+        :return: list of optimizers
+        """
+        optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
+        self.optimizers = [optimizer]
+        return self.optimizers
+
+    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(
+            dataset=dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=True
+        )
+
+        return loader
+
+    @property
+    def tng_dataloader(self):
+        if self._tng_dataloader is None:
+            try:
+                self._tng_dataloader = self.__dataloader(train=True)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._tng_dataloader
+
+    @property
+    def val_dataloader(self):
+        if self._val_dataloader is None:
+            try:
+                self._val_dataloader = self.__dataloader(train=False)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._val_dataloader
+
+    @property
+    def test_dataloader(self):
+        if self._test_dataloader is None:
+            try:
+                self._test_dataloader = self.__dataloader(train=False)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._test_dataloader
+
+    @staticmethod
+    def add_model_specific_args(parent_parser, root_dir):
+        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+
+        # param overwrites
+        # parser.set_defaults(gradient_clip=5.0)
+
+        # 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
+
+        # 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],
+                        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)
+        return parser

docs/source/examples/fully_featured_trainer.py

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

pyproject.toml

@@ -0,0 +1,5 @@
+[build-system]
+requires = [
+    "setuptools",
+    "wheel",
+]

pytorch_lightning/models/model_examples/bilstm.py

@@ -1,167 +0,0 @@
-import torch.nn as nn
-import numpy as np
-
-from test_tube import HyperOptArgumentParser
-import torch
-from torch.autograd import Variable
-from sklearn.metrics import confusion_matrix, f1_score
-from torch.nn import functional as F
-
-
-class BiLSTMPack(nn.Module):
-    """
-    Sample model to show how to define a template
-    """
-    def __init__(self, hparams):
-        # init superclass
-        super(BiLSTMPack, self).__init__(hparams)
-
-        self.hidden = None
-
-        # trigger tag building
-        self.ner_tagset = {'O': 0, 'I-Bio': 1}
-        self.nb_tags = len(self.ner_tagset)
-
-        # build model
-        print('building model...')
-        if hparams.model_load_weights_path is None:
-            self.__build_model()
-            print('model built')
-        else:
-            self = BiLSTMPack.load(hparams.model_load_weights_path, hparams.on_gpu, hparams)
-            print('model loaded from: {}'.format(hparams.model_load_weights_path))
-
-    def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
-        # design the number of final units
-        self.output_dim = self.hparams.nb_lstm_units
-
-        # when it's bidirectional our weights double
-        if self.hparams.bidirectional:
-            self.output_dim *= 2
-
-        # total number of words
-        total_words = len(self.tng_dataloader.dataset.words_token_to_idx)
-
-        # word embeddings
-        self.word_embedding = nn.Embedding(
-            num_embeddings=total_words + 1,
-            embedding_dim=self.hparams.embedding_dim,
-            padding_idx=0
-        )
-
-        # design the LSTM
-        self.lstm = nn.LSTM(
-            self.hparams.embedding_dim,
-            self.hparams.nb_lstm_units,
-            num_layers=self.hparams.nb_lstm_layers,
-            bidirectional=self.hparams.bidirectional,
-            dropout=self.hparams.drop_prob,
-            batch_first=True,
-        )
-
-        # map to tag space
-        self.fc_out = nn.Linear(self.output_dim, self.out_dim)
-        self.hidden_to_tag = nn.Linear(self.output_dim, self.nb_tags)
-
-
-    def init_hidden(self, batch_size):
-
-        # the weights are of the form (nb_layers * 2 if bidirectional, batch_size, nb_lstm_units)
-        mult = 2 if self.hparams.bidirectional else 1
-        hidden_a = torch.randn(self.hparams.nb_layers * mult, batch_size, self.nb_rnn_units)
-        hidden_b = torch.randn(self.hparams.nb_layers * mult, batch_size, self.nb_rnn_units)
-
-        if self.hparams.on_gpu:
-            hidden_a = hidden_a.cuda()
-            hidden_b = hidden_b.cuda()
-
-        hidden_a = Variable(hidden_a)
-        hidden_b = Variable(hidden_b)
-
-        return (hidden_a, hidden_b)
-
-    def forward(self, model_in):
-        # layout data (expand it, etc...)
-        # x = sequences
-        x, seq_lengths = model_in
-        batch_size, seq_len = x.size()
-
-        # reset RNN hidden state
-        self.hidden = self.init_hidden(batch_size)
-
-        # embed
-        x = self.word_embedding(x)
-
-        # run through rnn using packed sequences
-        x = torch.nn.utils.rnn.pack_padded_sequence(x, seq_lengths, batch_first=True)
-        x, self.hidden = self.lstm(x, self.hidden)
-        x, _ = torch.nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
-
-        # if asked for only last state, use the h_n which is the same as out(t=n)
-        if not self.return_sequence:
-            # pull out hidden states
-            # h_n = (nb_directions * nb_layers, batch_size, emb_size)
-            nb_directions = 2 if self.bidirectional else 1
-            (h_n, _) = self.hidden
-
-            # reshape to make indexing easier
-            # forward = 0, backward = 1 (of nb_directions)
-            h_n = h_n.view(self.nb_layers, nb_directions, batch_size, self.nb_rnn_units)
-
-            # pull out last forward
-            forward_h_n = h_n[-1, 0, :, :]
-            x = forward_h_n
-
-            # if bidirectional, also pull out the last hidden of backward network
-            if self.bidirectional:
-                backward_h_n = h_n[-1, 1, :, :]
-                x = torch.cat([forward_h_n, backward_h_n], dim=1)
-
-        # project to tag space
-        x = x.contiguous()
-        x = x.view(-1, self.output_dim)
-        x = self.hidden_to_tag(x)
-
-        return x
-
-    def loss(self, model_out):
-        # cross entropy loss
-        logits, y = model_out
-        y, y_lens = y
-
-        # flatten y and logits
-        y = y.view(-1)
-        logits = logits.view(-1, self.nb_tags)
-
-        # calculate a mask to remove padding tokens
-        mask = (y >= 0).float()
-
-        # count how many tokens we have
-        num_tokens = int(torch.sum(mask).data[0])
-
-        # pick the correct values and mask out
-        logits = logits[range(logits.shape[0]), y] * mask
-
-        # compute the ce loss
-        ce_loss = -torch.sum(logits)/num_tokens
-
-        return ce_loss
-
-    def pull_out_last_embedding(self, x, seq_lengths, batch_size, on_gpu):
-        # grab only the last activations from the non-padded ouput
-        x_last = torch.zeros([batch_size, 1, x.size(-1)])
-        for i, seq_len in enumerate(seq_lengths):
-            x_last[i, :, :] = x[i, seq_len-1, :]
-
-        # put on gpu when requested
-        if on_gpu:
-            x_last = x_last.cuda()
-
-        # turn into torch var
-        x_last = Variable(x_last)
-
-        return x_last

pytorch_lightning/models/trainer.py

@@ -1,31 +1,55 @@
 import torch
 import tqdm
 import numpy as np
-from research_lib.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module.memory import get_gpu_memory_map
 import traceback
-from research_lib.root_module.model_saving import TrainerIO
+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):
 
     def __init__(self,
                  experiment,
-                 cluster,
                  checkpoint_callback, early_stop_callback,
+                 cluster=None,
                  process_position=0,
                  current_gpu_name=0,
-                 on_gpu=False,
+                 gpus=None,
                  enable_tqdm=True,
-                 overfit_pct=None,
+                 overfit_pct=0.0,
                  track_grad_norm=-1,
                  check_val_every_n_epoch=1,
                  fast_dev_run=False,
-                 accumulate_grad_batches=False,
+                 accumulate_grad_batches=1,
                  enable_early_stop=True, max_nb_epochs=5, min_nb_epochs=1,
                  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
@@ -86,7 +119,7 @@ class Trainer(TrainerIO):
             self.test_percent_check = overfit_pct
 
     def __is_function_implemented(self, f_name):
-        f_op = getattr(self, f_name, None)
+        f_op = getattr(self.model, f_name, None)
         return callable(f_op)
 
     @property
@@ -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,18 +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:
@@ -217,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)
@@ -226,11 +278,13 @@ class Trainer(TrainerIO):
         self.experiment.save()
 
         # enable cluster checkpointing
-        self.enable_auto_hpc_walltime_manager()
+        if self.cluster is not None:
+            self.enable_auto_hpc_walltime_manager()
 
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
+        self.model = model
         self.__train()
 
     def __train(self):
@@ -240,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
@@ -257,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
@@ -268,25 +326,29 @@ class Trainer(TrainerIO):
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
-                self.__run_tng_batch(data_batch)
+                batch_result = self.__run_tng_batch(data_batch, batch_nb)
+                early_stop_epoch = batch_result == -1
 
                 # ---------------
                 # RUN VAL STEP
                 # ---------------
                 is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
-                if self.fast_dev_run or is_val_check_batch:
+                if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
                     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:
@@ -295,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
@@ -304,11 +368,17 @@ 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:
+                    break
 
             # 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:
@@ -320,24 +390,48 @@ class Trainer(TrainerIO):
                 if stop:
                     return
 
-    def __run_tng_batch(self, data_batch):
+
+    def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
-            return
+            return 0
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            self.model.on_batch_start()
+            model = self.model.module if self.data_parallel else self.model
+            response = model.on_batch_start(data_batch)
+
+            if response == -1:
+                return -1
 
         if self.enable_tqdm:
             self.prog_bar.update(1)
 
         # 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
@@ -368,6 +462,8 @@ class Trainer(TrainerIO):
         if self.__is_function_implemented('on_batch_end'):
             self.model.on_batch_end()
 
+        return 0
+
     def __run_validation(self):
         # decide if can check epochs
         can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0

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/hooks.py

@@ -1,7 +1,7 @@
 import torch
 
 class ModelHooks(torch.nn.Module):
-    def on_batch_start(self):
+    def on_batch_start(self, data_batch):
         pass
 
     def on_batch_end(self):
@@ -18,3 +18,4 @@ class ModelHooks(torch.nn.Module):
 
     def on_post_performance_check(self):
         pass
+

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

@@ -14,7 +14,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
     def __init__(self, hparams):
         super(RootModule, self).__init__()
         self.hparams = hparams
-        self.on_gpu = hparams.on_gpu
+
         self.dtype = torch.FloatTensor
         self.exp_save_path = None
         self.current_epoch = 0
@@ -24,6 +24,15 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.overfit = hparams.overfit
         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
+        try:
+            self.on_gpu = hparams.on_gpu
+        except Exception as e:
+            pass
 
         # computed vars for the dataloaders
         self._tng_dataloader = None
@@ -32,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):
         """
@@ -44,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:
@@ -60,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:
@@ -143,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/trainer_main.py

@@ -9,7 +9,6 @@ from pytorch_lightning.utils.arg_parse import add_default_args
 from time import sleep
 
 from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-
 SEED = 2334
 torch.manual_seed(SEED)
 np.random.seed(SEED)

pytorch_lightning/utils/arg_parse.py

@@ -1,4 +1,6 @@
-def add_default_args(parser, root_dir, possible_model_names, rand_seed):
+import pdb
+
+def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None):
 
     # tng, test, val check intervals
     parser.add_argument('--eval_test_set', dest='eval_test_set', action='store_true', help='true = run test set also')
@@ -32,7 +34,9 @@ def add_default_args(parser, root_dir, possible_model_names, rand_seed):
 
     # model paths
     parser.add_argument('--model_load_weights_path', default=None, type=str)
-    parser.add_argument('--model_name', default='', help=','.join(possible_model_names))
+
+    if possible_model_names is not None:
+        parser.add_argument('--model_name', default='', help=','.join(possible_model_names))
 
     # test_tube settings
     parser.add_argument('-en', '--tt_name', default='r_lib_')
@@ -43,10 +47,13 @@ def add_default_args(parser, root_dir, possible_model_names, rand_seed):
     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')
@@ -58,10 +65,12 @@ def add_default_args(parser, root_dir, possible_model_names, rand_seed):
     parser.add_argument('--overfit', default=-1, type=float, help='% of dataset to use with this option. float, or -1 for none')
 
     # debug args
-    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')
+    if rand_seed is not None:
+        parser.add_argument('--random_seed', default=rand_seed, type=int)
+
+    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/embeddings.py

@@ -13,11 +13,7 @@ class PretrainedEmbedding(torch.nn.Embedding):
         >>> emb = PretrainedEmbedding(embedding_path='glove.840B.300d.txt',embedding_dim=300, task_vocab={'hello': 1, 'world': 2})
         >>> data = torch.Tensor([[0, 1], [0, 2]]).long()
         >>> embedded = emb(data)
-        tensor([[[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
-         [ 0.2523,  0.1018, -0.6748,  ...,  0.1787, -0.5192,  0.3359]],
-
-        [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
-         [-0.0067,  0.2224,  0.2771,  ...,  0.0594,  0.0014,  0.0987]]])
+        
 
 
         :param embedding_path:
@@ -37,7 +33,8 @@ class PretrainedEmbedding(torch.nn.Embedding):
         self.weight = new_emb.weight
 
         # apply freeze
-        self.weight.requires_grad = not freeze
+        should_freeze = not freeze
+        self.weight.requires_grad = should_freeze
 
     def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
         """
@@ -97,11 +94,11 @@ class PretrainedEmbedding(torch.nn.Embedding):
 
 if __name__ == '__main__':
     emb = PretrainedEmbedding(
-        embedding_path='/Users/waf/Developer/NGV/research-fermat/fermat/.vector_cache/glove.840B.300d.txt',
+        embedding_path='/Users/waf/Developer',
         embedding_dim=300,
         task_vocab={'hello': 1, 'world': 2}
     )
 
     data = torch.Tensor([[0, 1], [0, 2]]).long()
     embedded = emb(data)
-    print(embedded)
+    print(embedded)

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.cfg

@@ -0,0 +1,21 @@
+[tool:pytest]
+norecursedirs =
+    .git
+    dist
+    build
+python_files =
+    test_*.py
+doctest_plus = disabled
+addopts = --strict
+markers =
+    slow
+    remote_data
+    filterwarnings
+
+[pycodestyle]
+ignore = E731,W504
+max-line-length = 120
+
+[flake8]
+ignore = E731,W504,F401,F841
+max-line-length = 120

setup.py

@@ -2,12 +2,28 @@
 
 from setuptools import setup, find_packages
 
-setup(name='pytorch-lightning',
-      version='0.0.1',
-      description='Rapid research framework',
-      author='',
-      author_email='',
-      url='https://github.com/williamFalcon/pytorch-lightning',
-      install_requires=[],
-      packages=find_packages()
-      )
+# https://packaging.python.org/guides/single-sourcing-package-version/
+
+# http://blog.ionelmc.ro/2014/05/25/python-packaging/
+setup(
+    name="pytorch-lightning",
+    version='0.11',
+    description="The Keras for ML researchers using PyTorch",
+    author="William Falcon",
+    author_email="waf2107@columbia.edu",
+    url="https://github.com/williamFalcon/pytorch-lightning",
+    download_url="https://github.com/williamFalcon/pytorch-lightning",
+    license="MIT",
+    keywords=["deep learning", "pytorch", "AI"],
+    python_requires=">=3.5",
+    install_requires=[
+        "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,
+)

tests/README.md

@@ -1,65 +0,0 @@
-# Testing setup
-
-## A. Enable CircleCI for your project
-1. Integrate CircleCI by clicking "Set up Project" at [this link](https://circleci.com/add-projects/gh/NextGenVest).
-
-## B. Add your own tests
-1. In the /tests, emulate exactly the folder structure for your module found under /bot_seed
-2. To create a test for file ```/bot_seed/folder/example.py```:
-    - create the file ```/tests/folder/example_test.py```
-    - notice the **_test**
-    - notice the mirror path under **/tests**
-
-3. Your ```example_test.py``` file should have these main components
-
-```python
-# example.py
-
-def function_i_want_to_test(x):
-    return x*2
-
-def square(x):
-    return x*x
-
-```
-
-```python
-# example_test.py
-
-import pytest
-
-# do whatever imports you need
-from app.bot_seed.folder.example import function_i_want_to_test, square
-
-def test_function_i_want_to_test():
-    answer = function_i_want_to_test(4)
-    assert answer == 8
-
-# -----------------------------------
-# Your function must start with test_
-# -----------------------------------
-def test_square():
-    answer = square(3)
-    assert answer == 9
-
-# -----------------------------------
-# boilerplate (link this file to pytest)
-# -----------------------------------
-if __name__ == '__main__':
-    pytest.main([__file__])
-```
-
-## C. Add build passing badge
-1. Create a CircleCI status token:
-    - Go here: https://circleci.com/gh/NextGenVest/your-project-name/edit#api
-    - Click create token
-    - Select status
-    - Type "badge status"
-
-2. Get a copy of the markdown code:
-    - Go here: https://circleci.com/gh/NextGenVest/your-project-name/edit#badges
-    - Select master
-    - Select "badge status" token
-    - Select image URL
-    - Copy the image url link and change the html at the top of the root README.md file for your project
-    

tests/research_proj/sample_model_template/model_template_test.py

@@ -1,13 +0,0 @@
-import pytest
-
-"""
-Example test to show how to add a test for anything in the project.  
-Look at the README for more instructions   
-"""
-
-
-def test_cube():
-    assert 27 == 27
-
-if __name__ == '__main__':
-    pytest.main([__file__])