fix install

* smurf ethics

* smurf ethics

* removed auto ddp fix

* removed auto ddp fix

* removed auto ddp fix

* removed auto ddp fix

* removed auto ddp fix

* removed auto ddp fix

.codecov.yml

@@ -1,4 +1,7 @@
-#see https://github.com/codecov/support/wiki/Codecov-Yaml
+# see https://docs.codecov.io/docs/codecov-yaml
+# Validation check:
+# $ curl --data-binary @.codecov.yml https://codecov.io/validate
+
 codecov:
   notify:
     require_ci_to_pass: yes
@@ -41,4 +44,3 @@ comment:
   behavior: default  # update if exists else create new
   # branches: *
 
-

.github/BECOMING_A_CORE_CONTRIBUTOR.md

@@ -0,0 +1,59 @@
+# How to become a core contributor
+
+Thanks for your interest in joining the Lightning team! We’re a rapidly growing project which is poised to become the go-to framework for DL researchers! 
+We're currently recruiting for a team of 5 core maintainers. 
+
+As a core maintainer you will have a strong say in the direction of the project. Big changes will require a majority of maintainers to agree.
+
+### Code of conduct  
+First and foremost, you'll be evaluated against [these core values](https://github.com/williamFalcon/pytorch-lightning/blob/master/.github/CONTRIBUTING.md). Any code we commit or feature we add needs to align with those core values.
+
+### The bar for joining the team   
+Lightning is being used to solve really hard problems at the top AI labs in the world. As such, the bar for adding team members is extremely high. Candidates must have solid engineering skills, have a good eye for user experience, and must be a power user of Lightning and PyTorch.
+
+With that said, the Lightning team will be diverse and a reflection of an inclusive AI community. You don't have to be an engineer to conntribute! Scientists with great usability intuition and PyTorch ninja skills are welcomed!
+
+### Responsibilities:
+The responsibilities mainly revolve around 3 things.
+
+#### Github issues
+- Here we want to help users have an amazing experience. These range from questions from new people getting into DL to questions from researchers about doing something esoteric with Lightning
+Often, these issues require some sort of bug fix, document clarification or new functionality to be scoped out.
+
+- To become a core member you must resolve at least 10 Github issues which align with the API design goals for Lightning. By the end of these 10 issues I should feel comfortable in the way you answer user questions
+Pleasant/helpful tone.
+
+- Can abstract from that issue or bug into functionality that might solve other related issues or makes the platform more flexible.
+
+- Don’t make users feel like they don’t know what they’re doing. We’re here to help and to make everyone’s experience delightful.
+
+#### Pull requests
+
+- Here we need to ensure the code that enters Lightning is high quality. For each PR we need to:
+- Make sure code coverage does not decrease
+- Documents are updated
+- Code is elegant and simple
+- Code is NOT overly engineered or hard to read
+- Ask yourself, could a non-engineer understand what’s happening here?
+- Make sure new tests are written
+- Is this NECESSARY for Lightning? There are some PRs which are just purely about adding engineering complexity which have no place in Lightning. 
+Guidance
+- Some other PRs are for people who are wanting to get involved and add something unnecessary. We do want their help though! So don’t approve the PR, but direct them to a Github issue that they might be interested in helping with instead!
+- To be considered for core contributor, please review 10 PRs and help the authors land it on master. Once you've finished the review, ping me 
+for a sanity check. At the end of 10 PRs if your PR reviews are inline with expectations described above, then you can merge PRs on your own going forward,
+otherwise we'll do a few more until we're both comfortable :)
+
+#### Project directions
+There are some big decisions which the project must make. For these I expect core contributors to have something meaningful to add if it’s their area of expertise.
+
+#### Diversity
+Lightning should reflect the broader community it serves. As such we should have scientists/researchers from 
+different fields contributing!   
+
+The first 5 core contributors will fit this profile. Thus if you overlap strongly with experiences and expertise as someone else on the team, you might have to wait until the next set of contributors are added.
+
+#### Summary: Requirements to apply
+- Solve 10 Github issues. The goal is to be inline with expectations for solving issues by the last one so you can do them on your own. If not, I might ask you to solve a few more specific ones.
+- Do 10 PR reviews. The goal is to be inline with expectations for solving issues by the last one so you can do them on your own. If not, I might ask you to solve a few more specific ones.
+
+If you want to be considered, ping me on gitter and start [tracking your progress here](https://docs.google.com/spreadsheets/d/15D58gp8DvI0Z6qbbYVRuaWioiwzafcP58-UlbuO_CMU/edit?usp=sharing).

.github/CODE_OF_CONDUCT.md

@@ -0,0 +1,76 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+ advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+ address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+ professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies both within project spaces and in public spaces
+when an individual is representing the project or its community. Examples of
+representing a project or community include using an official project e-mail
+address, posting via an official social media account, or acting as an appointed
+representative at an online or offline event. Representation of a project may be
+further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at waf2107@columbia.edu. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see
+https://www.contributor-covenant.org/faq

.github/CONTRIBUTING.md

@@ -1,10 +1,14 @@
 # Contributing    
-Welcome to the PyTorch Lightning community! We're building the most advanced research platform on the planet to implement the latest, best practices that the amazing PyTorch team rolls out!  
+Welcome to the PyTorch Lightning community! We're building the most advanced research platform on the planet to implement the latest, best practices that the amazing PyTorch team rolls out!   
+
+## Main Core Value: One less thing to remember    
+Simplify the API as much as possible from the user perspective. Any additions or improvements should minimize things the user needs to remember.   
+
+For example: One benefit of the validation_step is that the user doesn't have to remember to set the model to .eval(). This avoids all sorts of subtle errors the user could make.  
 
 ## Lightning Design Principles   
 We encourage all sorts of contributions you're interested in adding! When coding for lightning, please follow these principles.    
-
-#### No PyTorch interference   
+#### No PyTorch Interference   
 We don't want to add any abstractions on top of pure PyTorch. This gives researchers all the control they need without having to learn yet another framework.    
 
 #### Simple Internal Code    
@@ -16,17 +20,25 @@ There are 1,000 ways to do something. However, something eventually becomes stan
 When something becomes a best practice, we add it to the framework. This likely looks like code in utils or in the model file that everyone keeps adding over and over again across projects. When this happens, bring that code inside the trainer and add a flag for it.
 
 #### Simple External API    
-What makes sense to you may not make sense to others. Create an issue with an API change suggestion and validate that it makes sense for others. Treat code changes how you treat a startup: validate that it's a needed feature, then add if it makes sense for many people.    
+What makes sense to you may not make sense to others. Create an issue with an API change suggestion and validate that it makes sense for others. Treat code changes how you treat a startup: validate that it's a needed feature, then add if it makes sense for many people. 
+
+#### Backward-compatible API   
+We all hate updating our deep learning packages because we don't want to refactor a bunch of stuff. In Lightning, we make sure every change we make which could break an API is backwards compatible with good deprecation warnings.
+
+You shouldn't be afraid to upgrade Lightning :)
 
 #### Gain User Trust    
 As a researcher you can't have any part of your code going wrong. So, make thorough tests that ensure an implementation of a new trick or subbtle change is correct.    
 
-## Contribution types    
+#### Interoperability  
+Have a favorite feature from other libraries like fast.ai or transformers? Those should just work with lightning as well. Grab your favorite model or learning rate scheduler from your favorite library and run it in Lightning.
+
+## Contribution Types    
 Currently looking for help implementing new features or adding bug fixes. 
 
 A lot of good work has already been done in project mechanics (requirements.txt, setup.py, pep8, badges, ci, etc...) we're in a good state there thanks to all the early contributors (even pre-beta release)!   
 
-## Bug fixes:  
+## Bug Fixes:  
 1. Submit a github issue.   
 2. Fix it.  
 3. Submit a PR! 
@@ -35,3 +47,7 @@ A lot of good work has already been done in project mechanics (requirements.txt,
 1. Submit a github issue.   
 2. We'll agree on the feature scope.     
 3. Submit a PR! (with updated docs and tests 🙃).   
+
+## Coding Styleguide
+1. Test the code with flake8.
+2. Use f-strings.

.github/ISSUE_TEMPLATE/bug_report.md

@@ -7,6 +7,10 @@ assignees: ''
 
 ---
 
+### Common bugs:
+1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/williamFalcon/pytorch-lightning/issues/79).    
+2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/williamFalcon/pytorch-lightning#faq)    
+
 **Describe the bug**
 A clear and concise description of what the bug is.
 
@@ -28,11 +32,5 @@ If applicable, add screenshots to help explain your problem.
  - Browser [e.g. chrome, safari]
  - Version [e.g. 22]
 
-**Smartphone (please complete the following information):**
- - Device: [e.g. iPhone6]
- - OS: [e.g. iOS8.1]
- - Browser [e.g. stock browser, safari]
- - Version [e.g. 22]
-
 **Additional context**
 Add any other context about the problem here.

.github/ISSUE_TEMPLATE/how-to-question.md

@@ -0,0 +1,26 @@
+---
+name: How to question
+about: Asking how-to questions
+title: ''
+labels: question
+assignees: ''
+
+---
+
+### Before asking:   
+1. search the issues.   
+2. search the docs.    
+
+If you still can't find what you need:     
+#### What is your question?    
+
+#### Code    
+Please paste a code snippet if your question requires it!   
+
+#### What have you tried?    
+
+#### What's your environment?   
+- conda version (no venv)   
+- PyTorch version   
+- Lightning version   
+- Test-tube version

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,16 @@
+# Before submitting
+
+- Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
+- Did you read the [contributor guideline](https://github.com/williamFalcon/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)?
+- Did you make sure to update the docs?   
+- Did you write any new necessary tests?  
+
+## What does this PR do?
+Fixes # (issue).
+
+## PR review    
+Anyone in the community is free to review the PR once the tests have passed.     
+If we didn't discuss your PR in Github issues there's a high chance it will not be merged.   
+
+## Did you have fun?
+Make sure you had fun coding 🙃

.gitignore

@@ -11,6 +11,9 @@ pip-wheel-metadata/
 test_tube_exp/
 tests/tests_tt_dir/
 tests/save_dir
+default/
+lightning_logs/
+tests/tests/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

.run_local_tests.sh

@@ -0,0 +1,8 @@
+# use this to run tests
+rm -rf _ckpt_*
+rm -rf tests/save_dir*
+rm -rf tests/mlruns_*
+rm -rf tests/tests/*
+rm -rf lightning_logs
+coverage run --source pytorch_lightning -m py.test pytorch_lightning tests pl_examples -v --doctest-modules
+coverage report -m

.travis.yml

@@ -8,8 +8,6 @@
 # this file is *not* meant to cover or endorse the use of travis, but rather to
 # help confirm pull requests to this project.
 
-dist: xenial  # Ubuntu 16.04
-
 env:
   global:
     - DISPLAY=""
@@ -18,23 +16,48 @@ language: python
 
 matrix:
   include:
-    - python: 3.6
+    - os: linux
+      dist: xenial  # Ubuntu 16.04
+      python: 3.6
       env: TOXENV=py36
-    - python: 3.7
+    - os: linux
+      dist: bionic  # Ubuntu 18.04
+      python: 3.6
+      env: TOXENV=py36
+    - os: linux
+      dist: bionic  # Ubuntu 18.04
+      python: 3.7
       env: TOXENV=py37
+    - os: osx
+      osx_image: xcode9.4
+      language: generic
+      env: TOXENV=py36
+      addons:
+        homebrew:
+          # update: true
+          packages: python3.6
+      before_install:
+        - pip3 install virtualenv
+        - virtualenv -p python3 ~/venv
+        - source ~/venv/bin/activate
+#    - os: windows
+#      language: minimal
+#      before_install:
+#        - choco install python3
+#        - export PATH="/c/Python37:/c/Python37/Scripts:$PATH"
+#      env: TOXENV=py37
 
 # See http://docs.travis-ci.com/user/caching/#pip-cache
 cache: pip
 
 install:
-  - pip install -r requirements.txt
-  - pip install -r ./tests/requirements.txt
-  - pip --version ; pip list
+  - pip install future  # needed for `builtins`
+  - sudo pip install tox
 
 script:
   # integration
   - tox --sitepackages
-  - python setup.py install --dry-run
+  - pip install --editable .
 
 after_success:
   - coverage report

MANIFEST.in

@@ -13,11 +13,16 @@ include LICENSE
 exclude *.sh
 exclude *.toml
 exclude *.svg    
-recursive-include examples *.py
 recursive-include pytorch_lightning *.py
 
+# include examples
+recursive-include pl_examples *.py
+recursive-include pl_examples *.md
+recursive-include pl_examples *.sh
+
 # exclude tests from package
 recursive-exclude tests *
+recursive-exclude site *
 exclude tests
 
 # Exclude the documentation files

README.md

@@ -4,23 +4,25 @@
 
 # PyTorch Lightning
 
-**The PyTorch Keras for ML researchers. More control. Less boilerplate.**
+**The lightweight PyTorch wrapper for ML researchers. Scale your models. Write less boilerplate.**
 
 
 [![PyPI Status](https://badge.fury.io/py/pytorch-lightning.svg)](https://badge.fury.io/py/pytorch-lightning)
 [![PyPI Status](https://pepy.tech/badge/pytorch-lightning)](https://pepy.tech/project/pytorch-lightning)
 [![Build Status](https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master)](https://travis-ci.org/williamFalcon/pytorch-lightning)
-<!-- 
-removed until windows install issues resolved.
-[![Build status](https://ci.appveyor.com/api/projects/status/rum89d7hq8l1kfye?svg=true)](https://ci.appveyor.com/project/Borda/pytorch-lightning) -->
+[![Build status](https://ci.appveyor.com/api/projects/status/NEW-PROJECT-ID?svg=true)](https://ci.appveyor.com/project/williamFalcon/pytorch-lightning)
+[![Coverage](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/coverage.svg)](https://github.com/williamFalcon/pytorch-lightning/tree/master/tests#running-coverage)
+[![CodeFactor](https://www.codefactor.io/repository/github/borda/pytorch-lightning/badge)](https://www.codefactor.io/repository/github/borda/pytorch-lightning)    
+
+[![ReadTheDocs](https://readthedocs.org/projects/pytorch-lightning/badge/?version=latest)](https://pytorch-lightning.readthedocs.io/en/latest)
+[![Gitter](https://badges.gitter.im/PyTorch-Lightning/community.svg)](https://gitter.im/PyTorch-Lightning/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge)
+[![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/williamFalcon/pytorch-lightning/blob/master/LICENSE)
+[![Next Release](https://img.shields.io/badge/Next%20Release-Dec%206-<COLOR>.svg)](https://shields.io/)
+
 <!-- 
 removed until codecov badge isn't empy. likely a config error showing nothing on master.
 [![codecov](https://codecov.io/gh/Borda/pytorch-lightning/branch/master/graph/badge.svg)](https://codecov.io/gh/Borda/pytorch-lightning)
 -->
-[![Coverage](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/coverage.svg)](https://github.com/williamFalcon/pytorch-lightning/tree/master/tests#running-coverage)
-[![CodeFactor](https://www.codefactor.io/repository/github/borda/pytorch-lightning/badge)](https://www.codefactor.io/repository/github/borda/pytorch-lightning)
-[![ReadTheDocs](https://readthedocs.org/projects/pytorch-lightning/badge/?version=latest)](https://pytorch-lightning.readthedocs.io/en/latest)
-[![license](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/williamFalcon/pytorch-lightning/blob/master/LICENSE)
 
 </div>
 
@@ -33,34 +35,42 @@ pip install pytorch-lightning
 **[View the docs here](https://williamfalcon.github.io/pytorch-lightning/)**
 
 ## What is it?  
-Lightning is a very lightweight wrapper on PyTorch. This means you don't have to learn a new library. It defers core training and validation logic to you and automates the rest. It guarantees tested, correct, modern best practices for the automated parts.
+Lightning is a very lightweight wrapper on PyTorch. This means you don't have to learn a new library. To use Lightning, simply refactor your research code into the [LightningModule](https://github.com/williamFalcon/pytorch-lightning#how-do-i-do-use-it) format and Lightning will automate the rest. Lightning guarantees tested, correct, modern best practices for the automated parts.
 
+## Starting a new project?   
+[Use our seed-project aimed at reproducibility!](https://github.com/williamFalcon/pytorch-lightning-conference-seed)     
 
 ## Why do I want to use lightning?
-When starting a new project the last thing you want to do is recode a training loop, multi-cluster training, 16-bit precision, early-stopping, model loading/saving, when to validate, etc... You're likely to spend a long time ironing out all the bugs without even getting to the core of your research.
+Every research project starts the same, a model, a training loop, validation loop, etc. As your research advances, you're likely to need distributed training, 16-bit precision, checkpointing, gradient accumulation, etc.   
 
-With lightning, you guarantee those parts of your code work so you can focus on what the meat of the research: The data and the training/validation loop logic. 
+Lightning sets up all the boilerplate state-of-the-art training for you so you can focus on the research.   
 
-Don't worry about training on multiple gpus or speeding up your code, lightning will do that for you!
-
----   
+---
+ 
 ## README Table of Contents        
 - [How do I use it](https://github.com/williamFalcon/pytorch-lightning#how-do-i-do-use-it)     
 - [What lightning automates](https://github.com/williamFalcon/pytorch-lightning#what-does-lightning-control-for-me)    
 - [Tensorboard integration](https://github.com/williamFalcon/pytorch-lightning#tensorboard)    
 - [Lightning features](https://github.com/williamFalcon/pytorch-lightning#lightning-automates-all-of-the-following-each-is-also-configurable)    
-- [Demos](https://github.com/williamFalcon/pytorch-lightning#demo)    
+- [Examples](https://github.com/williamFalcon/pytorch-lightning#examples)    
 - [Tutorials](https://github.com/williamFalcon/pytorch-lightning#tutorials)
-- [Contributing](https://github.com/williamFalcon/pytorch-lightning/blob/master/CONTRIBUTING.md)    
+- [Contributing](https://github.com/williamFalcon/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)
 - [Bleeding edge install](https://github.com/williamFalcon/pytorch-lightning#bleeding-edge)   
-- [Lightning Design Principles](https://github.com/williamFalcon/pytorch-lightning#lightning-design-principles)
+- [Lightning Design Principles](https://github.com/williamFalcon/pytorch-lightning#lightning-design-principles)   
+- [Asking for help](https://github.com/williamFalcon/pytorch-lightning#asking-for-help)
 - [FAQ](https://github.com/williamFalcon/pytorch-lightning#faq)    
 
----   
+---
+
 ## How do I do use it?   
+Think about Lightning as refactoring your research code instead of using a new framework. The research code goes into a [LightningModule]((https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)) which you fit using a Trainer.    
+
+The LightningModule defines a *system* such as seq-2-seq, GAN, etc... It can ALSO define a simple classifier such as the example below.     
 
 To use lightning do 2 things:  
-1. [Define a LightningModel](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)         
+1. [Define a LightningModule](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)     
+
+**WARNING:** This syntax is for version 0.5.0+ where abbreviations were removed.    
 ```python
 import os
 import torch
@@ -71,104 +81,117 @@ import torchvision.transforms as transforms
 
 import pytorch_lightning as pl
 
-class CoolModel(pl.LightningModule):
+class CoolSystem(pl.LightningModule):
 
     def __init__(self):
-        super(CoolModel, self).__init__()
+        super(CoolSystem, self).__init__()
         # not the best model...
         self.l1 = torch.nn.Linear(28 * 28, 10)
 
     def forward(self, x):
         return torch.relu(self.l1(x.view(x.size(0), -1)))
 
-    def my_loss(self, y_hat, y):
-        return F.cross_entropy(y_hat, y)
-
     def training_step(self, batch, batch_nb):
+        # REQUIRED
         x, y = batch
         y_hat = self.forward(x)
-        return {'loss': self.my_loss(y_hat, y)}
+        loss = F.cross_entropy(y_hat, y)
+        tensorboard_logs = {'train_loss': loss}
+        return {'loss': loss, 'log': tensorboard_logs}
 
     def validation_step(self, batch, batch_nb):
+        # OPTIONAL
         x, y = batch
         y_hat = self.forward(x)
-        return {'val_loss': self.my_loss(y_hat, y)}
+        return {'val_loss': F.cross_entropy(y_hat, y)}
 
     def validation_end(self, outputs):
+        # OPTIONAL
         avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
-        return {'avg_val_loss': avg_loss}
+        tensorboard_logs = {'val_loss': avg_loss}
+        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}
 
     def configure_optimizers(self):
-        return [torch.optim.Adam(self.parameters(), lr=0.02)]
+        # REQUIRED
+        # can return multiple optimizers and learning_rate schedulers
+        # (LBFGS it is automatically supported, no need for closure function)
+        return torch.optim.Adam(self.parameters(), lr=0.02)
 
     @pl.data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
+        # REQUIRED
         return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @pl.data_loader
     def val_dataloader(self):
+        # OPTIONAL
         return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @pl.data_loader
     def test_dataloader(self):
-        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
+        # OPTIONAL
+        return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
-
 2. Fit with a [trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)    
 ```python
 from pytorch_lightning import Trainer
 
-model = CoolModel()
+model = CoolSystem()
 
 # most basic trainer, uses good defaults
 trainer = Trainer()    
 trainer.fit(model)   
 ```     
 
-Or with tensorboard logger and some options turned on such as multi-gpu, etc...    
+Trainer sets up a tensorboard logger, early stopping and checkpointing by default (you can modify all of them or
+use something other than tensorboard).   
+
+Here are more advanced examples
 ```python   
-from test_tube import Experiment    
-
-# PyTorch summarywriter with a few bells and whistles    
-exp = Experiment(save_dir=os.getcwd())
-
 # train on cpu using only 10% of the data (for demo purposes)
-# pass in experi
-trainer = Trainer(experiment=exp, max_nb_epochs=1, train_percent_check=0.1)
+trainer = Trainer(max_nb_epochs=1, train_percent_check=0.1)
 
-# train on 4 gpus
-# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3])
+# train on 4 gpus (lightning chooses GPUs for you)
+# trainer = Trainer(max_nb_epochs=1, gpus=4, distributed_backend='ddp')  
+
+# train on 4 gpus (you choose GPUs)
+# trainer = Trainer(max_nb_epochs=1, gpus=[0, 1, 3, 7], distributed_backend='ddp')   
 
 # train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)
-# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], nb_gpu_nodes=4)
+# trainer = Trainer(max_nb_epochs=1, gpus=8, nb_gpu_nodes=4, distributed_backend='ddp')
 
 # train (1 epoch only here for demo)
 trainer.fit(model)
 
-# view tensorflow logs 
-print('View tensorboard logs by running\ntensorboard --logdir %s' % os.getcwd())
-print('and going to http://localhost:6006 on your browser')
+# view tensorboard logs 
+logging.info(f'View tensorboard logs by running\ntensorboard --logdir {os.getcwd()}')
+logging.info('and going to http://localhost:6006 on your browser')
 ```    
 
+When you're all done you can even run the test set separately.   
+```python
+trainer.test()
+```
+
 ## What does lightning control for me?   
 
 Everything in gray!    
 You define the blue parts using the LightningModule interface:  
 
-![Ouverview](./docs/source/_static/overview_flat.jpg)
+![Overview](./docs/source/_static/overview_flat.jpg)
 
-```{.python}
+```python
 # what to do in the training loop
-def training_step(self, data_batch, batch_nb):
+def training_step(self, batch, batch_nb):
 
 # what to do in the validation loop
-def validation_step(self, data_batch, batch_nb):
+def validation_step(self, batch, batch_nb):
 
 # how to aggregate validation_step outputs
 def validation_end(self, outputs):
 
 # and your dataloaders
-def tng_dataloader():
+def train_dataloader():
 def val_dataloader():
 def test_dataloader():
 ```
@@ -177,8 +200,8 @@ def test_dataloader():
 
 ```python
 # define what happens for training here
-def training_step(self, data_batch, batch_nb):
-    x, y = data_batch
+def training_step(self, batch, batch_nb):
+    x, y = batch
     
     # define your own forward and loss calculation
     hidden_states = self.encoder(x)
@@ -204,8 +227,8 @@ def training_step(self, data_batch, batch_nb):
 
 ```python
 # define what happens for validation here
-def validation_step(self, data_batch, batch_nb):    
-    x, y = data_batch
+def validation_step(self, batch, batch_nb):    
+    x, y = batch
     
     # or as basic as a CNN classification
     out = self.forward(x)
@@ -230,12 +253,13 @@ def validation_end(self, outputs):
 
     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
+    logs = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+    result = {'log': logs}
+    return result
 ```
    
 ## Tensorboard    
-Lightning is fully integrated with tensorboard.   
+Lightning is fully integrated with tensorboard, MLFlow and supports any logging module.   
 
 ![tensorboard-support](./docs/source/_static/tf_loss.png)
 
@@ -243,36 +267,21 @@ Lightning also adds a text column with all the hyperparameters for this experime
 
 ![tensorboard-support](./docs/source/_static/tf_tags.png)
 
-Simply note the path you set for the Experiment    
-``` {.python}   
-from test_tube import Experiment
-from pytorch-lightning import  Trainer
-
-exp = Experiment(save_dir='/some/path')
-trainer = Trainer(experiment=exp)
-...
-```   
-
-And run tensorboard from that dir   
-```bash
-tensorboard --logdir /some/path     
-```    
-
 ## Lightning automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
 
+#### Checkpointing    
 
-###### Checkpointing    
-
+- [Checkpoint callback](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
 - [Restoring training session](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#restoring-training-session)
 
-###### Computing cluster (SLURM)    
+#### Computing cluster (SLURM)    
 
 - [Running grid search on a cluster](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#running-grid-search-on-a-cluster) 
 - [Walltime auto-resubmit](https://williamfalcon.github.io/pytorch-lightning/Trainer/SLURM%20Managed%20Cluster#walltime-auto-resubmit)   
 
-###### Debugging  
+#### Debugging  
 
 - [Fast dev run](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#fast-dev-run)
 - [Inspect gradient norms](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#inspect-gradient-norms)
@@ -283,8 +292,9 @@ tensorboard --logdir /some/path
 - [Print input and output size of every module in system](https://williamfalcon.github.io/pytorch-lightning/LightningModule/properties/#example_input_array)
 
 
-###### Distributed training    
+#### Distributed training    
 
+- [Implement Your Own Distributed (DDP) training](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#init_ddp_connection)
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
 - [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
@@ -292,7 +302,7 @@ tensorboard --logdir /some/path
 - [Self-balancing architecture](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#self-balancing-architecture)
 
 
-###### Experiment Logging   
+#### Experiment Logging   
 
 - [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
@@ -302,18 +312,20 @@ tensorboard --logdir /some/path
 - [Snapshot code for a training run](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#snapshot-code-for-a-training-run) 
 - [Write logs file to csv every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#write-logs-file-to-csv-every-k-batches)
 
-###### Training loop    
+#### Training loop    
 
 - [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
 - [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
+- [Early stopping callback](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#early-stopping)
 - [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
 - [Gradient Clipping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
 - [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
 - [Learning rate scheduling](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)    
 - [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
 - [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
+- [Step optimizers at arbitrary intervals](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step)
 
-###### Validation loop    
+#### Validation loop    
 
 - [Check validation every n epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#check-validation-every-n-epochs)
 - [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
@@ -322,36 +334,38 @@ tensorboard --logdir /some/path
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
 - [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
 
+#### Testing loop  
+- [Run test set](https://williamfalcon.github.io/pytorch-lightning/Trainer/Testing%20loop/)  
 
-## Demo
-```bash
-# install lightning
-pip install pytorch-lightning
-
-# clone lightning for the demo
-git clone https://github.com/williamFalcon/pytorch-lightning.git
-cd pytorch-lightning
-cd examples/new_project_templates/
-
-# all of the following demos use the SAME model to show no modification needs to be made to your code
-
-# train on cpu 
-python single_cpu_template.py
-
-# train on multiple-gpus 
-python single_gpu_node_template.py --gpus "0,1"
-
-# train on 32 gpus on a cluster (run on a SLURM managed cluster)
-python multi_node_cluster_template.py --nb_gpu_nodes 4 --gpus '0,1,2,3,4,5,6,7'
-```
+## Examples   
+- [GAN](https://github.com/williamFalcon/pytorch-lightning/tree/master/pl_examples/domain_templates/gan.py)    
+- [MNIST](https://github.com/williamFalcon/pytorch-lightning/tree/master/pl_examples/basic_examples)      
+- [Other projects using Lightning](https://github.com/williamFalcon/pytorch-lightning/network/dependents?package_id=UGFja2FnZS0zNzE3NDU4OTM%3D)    
+- [Multi-node](https://github.com/williamFalcon/pytorch-lightning/tree/master/pl_examples/multi_node_examples)   
 
 ## Tutorials   
 - [Basic Lightning use](https://towardsdatascience.com/supercharge-your-ai-research-with-pytorch-lightning-337948a99eec)    
 - [9 key speed features in Pytorch-Lightning](https://towardsdatascience.com/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565)    
 - [SLURM, multi-node training with Lightning](https://towardsdatascience.com/trivial-multi-node-training-with-pytorch-lightning-ff75dfb809bd)     
 
+---    
+## Asking for help    
+Welcome to the Lightning community!   
+
+If you have any questions, feel free to:   
+1. [read the docs](https://williamfalcon.github.io/pytorch-lightning/).     
+2. [Search through the issues](https://github.com/williamFalcon/pytorch-lightning/issues?utf8=%E2%9C%93&q=my++question).      
+3. [Ask on stackoverflow](https://stackoverflow.com/questions/ask?guided=false) with the tag pytorch-lightning.   
+
+If no one replies to you quickly enough, feel free to post the stackoverflow link to our Gitter chat!   
+
+To chat with the rest of us visit our [gitter channel](https://gitter.im/PyTorch-Lightning/community)!     
+
 ---   
 ## FAQ    
+**How do I use Lightning for rapid research?**   
+[Here's a walk-through](https://williamfalcon.github.io/pytorch-lightning/)  
+
 **Why was Lightning created?**     
 Lightning has 3 goals in mind:
 1. Maximal flexibility while abstracting out the common boilerplate across research projects.   
@@ -370,8 +384,50 @@ Nope.
 **Are there plans to support virtualenv?**    
 Nope. Please use anaconda or miniconda.    
 
-## Bleeding edge
-If you can't wait for the next release, install the most up to date code with:  
+**Which PyTorch versions do you support?**    
+- **PyTorch 1.1.0**       
+    ```bash    
+    # install pytorch 1.1.0 using the official instructions   
+    
+    # install test-tube 0.6.7.6 which supports 1.1.0   
+    pip install test-tube==0.6.7.6   
+    
+    # install latest Lightning version without upgrading deps    
+    pip install -U --no-deps pytorch-lightning
+    ```     
+- **PyTorch 1.2.0, 1.3.0,**
+    Install via pip as normal   
+
+## Custom installation
+
+### Bleeding edge
+
+If you can't wait for the next release, install the most up to date code with:
+* using GIT (locally clone whole repo with full history)
+    ```bash
+    pip install git+https://github.com/williamFalcon/pytorch-lightning.git@master --upgrade
+    ```
+* using instant zip (last state of the repo without git history)
+    ```bash
+    pip install https://github.com/williamFalcon/pytorch-lightning/archive/master.zip --upgrade
+    ```
+
+### Any release installation
+
+You can also install any past release from this repository:
 ```bash
-pip install git+https://github.com/williamFalcon/pytorch-lightning.git@master --upgrade
+pip install https://github.com/williamFalcon/pytorch-lightning/archive/0.4.4.zip --upgrade
+```
+
+## Bibtex
+If you want to cite the framework feel free to use this (but only if you loved it 😊):
+```
+@misc{Falcon2019,
+  author = {Falcon, W.A.},
+  title = {PyTorch Lightning},
+  year = {2019},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/williamFalcon/pytorch-lightning}}
+}
 ```

appveyor.yml

@@ -45,9 +45,7 @@ install:
   # directly to master instead of just PR builds (or the converse).
   - SET PATH=%PYTHON%;%PYTHON%\\Scripts;%path%
   - pip install -U --user pip
-  - pip install "https://download.pytorch.org/whl/cu90/torch-1.1.0-cp%PIP_PYVER%-cp%PIP_PYVER%m-win_amd%PYTHON_ARCH%.whl"
-    pip install "https://download.pytorch.org/whl/cu90/torchvision-0.3.0-cp%PIP_PYVER%-cp%PIP_PYVER%m-win_amd%PYTHON_ARCH%.whl"
-  - pip install -r requirements.txt
+  - pip install -r requirements.txt -f https://download.pytorch.org/whl/torch_stable.html
   - pip install -r ./tests/requirements.txt
 
 # scripts to run before tests (working directory and environment changes are persisted from the previous steps such as "before_build")

docs/LightningModule/RequiredTrainerInterface.md

@@ -9,22 +9,22 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 
 **Required**:  
 
-- [training_step](RequiredTrainerInterface.md#training_step)   
-- [validation_step](RequiredTrainerInterface.md#validation_step)
-- [validation_end](RequiredTrainerInterface.md#validation_end)
-    
-- [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)
-
-- [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
-- [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
-- [test_dataloader](RequiredTrainerInterface.md#test_dataloader)
+- [training_step](RequiredTrainerInterface.md#training_step)      
+- [train_dataloader](RequiredTrainerInterface.md#train_dataloader)    
+- [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)    
 
 **Optional**:   
 
-- [on_save_checkpoint](RequiredTrainerInterface.md#on_save_checkpoint)
-- [on_load_checkpoint](RequiredTrainerInterface.md#on_load_checkpoint)
-- [update_tng_log_metrics](RequiredTrainerInterface.md#update_tng_log_metrics)
-- [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)
+- [training_end](RequiredTrainerInterface.md#training_end)      
+- [validation_step](RequiredTrainerInterface.md#validation_step)    
+- [validation_end](RequiredTrainerInterface.md#validation_end) 
+- [test_step](RequiredTrainerInterface.md#test_step)    
+- [test_end](RequiredTrainerInterface.md#test_end) 
+- [val_dataloader](RequiredTrainerInterface.md#val_dataloader)    
+- [test_dataloader](RequiredTrainerInterface.md#test_dataloader)    
+- [on_save_checkpoint](RequiredTrainerInterface.md#on_save_checkpoint)    
+- [on_load_checkpoint](RequiredTrainerInterface.md#on_load_checkpoint)    
+- [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)    
 
 ---
 ### Minimal example   
@@ -48,37 +48,53 @@ class CoolModel(pl.LightningModule):
     def forward(self, x):
         return torch.relu(self.l1(x.view(x.size(0), -1)))
 
-    def my_loss(self, y_hat, y):
-        return F.cross_entropy(y_hat, y)
-
     def training_step(self, batch, batch_nb):
+        # REQUIRED
         x, y = batch
         y_hat = self.forward(x)
-        return {'loss': self.my_loss(y_hat, y)}
+        return {'loss': F.cross_entropy(y_hat, y)}
 
     def validation_step(self, batch, batch_nb):
+        # OPTIONAL
         x, y = batch
         y_hat = self.forward(x)
-        return {'val_loss': self.my_loss(y_hat, y)}
+        return {'val_loss': F.cross_entropy(y_hat, y)}
 
     def validation_end(self, outputs):
+        # OPTIONAL
         avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
         return {'avg_val_loss': avg_loss}
 
+    def test_step(self, batch, batch_nb):
+        # OPTIONAL
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'test_loss': F.cross_entropy(y_hat, y)}
+
+    def test_end(self, outputs):
+        # OPTIONAL
+        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
+        return {'avg_test_loss': avg_loss}
+
     def configure_optimizers(self):
-        return [torch.optim.Adam(self.parameters(), lr=0.02)]
+        # REQUIRED
+        return torch.optim.Adam(self.parameters(), lr=0.02)
 
     @pl.data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
         return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @pl.data_loader
     def val_dataloader(self):
+        # OPTIONAL
+        # can also return a list of val dataloaders
         return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 
     @pl.data_loader
     def test_dataloader(self):
-        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
+        # OPTIONAL
+        # can also return a list of test dataloaders
+        return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)
 ```
 ---   
 ### How do these methods fit into the broader training?     
@@ -90,12 +106,12 @@ The LightningModule interface is on the right. Each method corresponds to a part
   </a>
 </p>   
 
----
+## Required Methods    
 
 ### training_step
 
 ``` {.python}
-def training_step(self, data_batch, batch_nb)
+def training_step(self, batch, batch_nb)
 ```
 
 In this step you'd normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something specific to your model.
@@ -104,7 +120,7 @@ In this step you'd normally do the forward pass and calculate the loss for a bat
 
 | Param  | description  |
 |---|---|
-|  data_batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch | The output of your dataloader. A tensor, tuple or list  |
 |  batch_nb | Integer displaying which batch this is  |
 
 **Return**   
@@ -114,61 +130,258 @@ Dictionary or OrderedDict
 | key  | value  | is required |
 |---|---|---|
 |  loss | tensor scalar  | Y |
-|  prog | Dict for progress bar display. Must have only tensors  | N |
+|  progress_bar | Dict for progress bar display. Must have only tensors  | N |
+|  log | Dict of metrics to add to logger. Must have only tensors (no images, etc)  | N |
 
 
 **Example**
 
 ``` {.python}
-def training_step(self, data_batch, batch_nb):
-    x, y, z = data_batch
+def training_step(self, batch, batch_nb):
+    x, y, z = batch
     
     # implement your own
     out = self.forward(x)
     loss = self.loss(out, x)
     
+    logger_logs = {'training_loss': loss} # optional (MUST ALL BE TENSORS)
+    
+    # if using TestTubeLogger or TensorboardLogger you can nest scalars
+    logger_logs = {'losses': logger_logs} # optional (MUST ALL BE TENSORS)
+    
     output = {
         'loss': loss, # required
-        'prog': {'tng_loss': loss, 'batch_nb': batch_nb} # optional
+        'progress_bar': {'training_loss': loss}, # optional (MUST ALL BE TENSORS)
+        'log': logger_logs
     }
     
     # return a dict
     return output
+```    
+
+If you define multiple optimizers, this step will also be called with an additional ```optimizer_idx``` param.    
+``` {.python}
+# Multiple optimizers (ie: GANs)     
+def training_step(self, batch, batch_nb, optimizer_idx):
+    if optimizer_idx == 0:
+        # do training_step with encoder
+    if optimizer_idx == 1:
+        # do training_step with decoder    
+```    
+
+If you add truncated back propagation through time you will also get an additional argument with the hidden states of the previous step.     
+``` {.python}
+# Truncated back-propagation through time   
+def training_step(self, batch, batch_nb, hiddens):
+    # hiddens are the hiddens from the previous truncated backprop step
+```    
+
+You can also return a -1 instead of a dict to stop the current loop. This is useful if you want to
+break out of the current training epoch early.
+
+---   
+### training_end
+
+``` {.python}
+def training_end(self, train_step_outputs)
+```
+In certain cases (dp, ddp2), you might want to use all outputs of every process to do something.
+For instance, if using negative samples, you could run a batch via dp and use ALL the outputs 
+for a single softmax across the full batch (ie: the denominator would use the full batch).
+
+In this case you should define training_end to perform those calculations.
+
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  outputs | What you return in training_step. 
+
+**Return**   
+
+Dictionary or OrderedDict   
+
+| key  | value  | is required |
+|---|---|---|
+|  loss | tensor scalar  | Y |
+|  progress_bar | Dict for progress bar display. Must have only tensors  | N |
+|  log | Dict of metrics to add to logger. Must have only tensors (no images, etc)  | N |
+
+
+**Example**
+
+``` {.python}
+# WITHOUT training_end
+# if used in DP or DDP2, this batch is 1/nb_gpus large
+def training_step(self, batch, batch_nb):
+    # batch is 1/nb_gpus big
+    x, y = batch
+    
+    out = self.forward(x)
+    loss = self.softmax(out)
+    loss = nce_loss(loss)
+    return {'loss': loss}
+
+# --------------
+# with training_end to do softmax over the full batch
+def training_step(self, batch, batch_nb):
+    # batch is 1/nb_gpus big
+    x, y = batch
+    
+    out = self.forward(x)
+    return {'out': out}
+
+def training_end(self, outputs):
+    # this out is now the full size of the batch
+    out = outputs['out']
+
+    # this softmax now uses the full batch size
+    loss = self.softmax(out)
+    loss = nce_loss(loss)
+    return {'loss': loss}
+```    
+
+If you define multiple optimizers, this step will also be called with an additional ```optimizer_idx``` param.    
+``` {.python}
+# Multiple optimizers (ie: GANs)     
+def training_step(self, batch, batch_nb, optimizer_idx):
+    if optimizer_idx == 0:
+        # do training_step with encoder
+    if optimizer_idx == 1:
+        # do training_step with decoder    
+```    
+
+If you add truncated back propagation through time you will also get an additional argument with the hidden states of the previous step.     
+``` {.python}
+# Truncated back-propagation through time   
+def training_step(self, batch, batch_nb, hiddens):
+    # hiddens are the hiddens from the previous truncated backprop step
+```    
+
+You can also return a -1 instead of a dict to stop the current loop. This is useful if you want to
+break out of the current training epoch early.
+
+--- 
+### train_dataloader 
+
+``` {.python}
+@pl.data_loader
+def train_dataloader(self)
+```
+Called by lightning during training loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.   
+If you want to change the data during every epoch DON'T use the data_loader decorator.
+
+##### Return
+PyTorch DataLoader
+
+**Example**
+
+``` {.python}
+@pl.data_loader
+def train_dataloader(self):
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+    dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform, download=True)
+    loader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=self.hparams.batch_size,
+        shuffle=True
+    )
+    return loader
 ```
 
----
+--- 
+### configure_optimizers 
+
+``` {.python}
+def configure_optimizers(self)
+```
+
+Set up as many optimizers and (optionally) learning rate schedulers as you need. Normally you'd need one. But in the case of GANs or something more esoteric you might have multiple. 
+Lightning will call .backward() and .step() on each one in every epoch.  If you use 16 bit precision it will also handle that.
+
+**Note:** If you use multiple optimizers, training_step will have an additional ```optimizer_idx``` parameter.    
+**Note 2:** If you use LBFGS lightning handles the closure function automatically for you.
+
+##### Return    
+Return any of these 3 options:   
+Single optimizer   
+List or Tuple - List of optimizers    
+Two lists - The first list has multiple optimizers, the second a list of learning-rate schedulers
+
+**Example**
+
+``` {.python}
+# most cases
+def configure_optimizers(self):
+    opt = Adam(self.parameters(), lr=0.01)
+    return opt
+
+# multiple optimizer case (eg: GAN)
+def configure_optimizers(self):
+    generator_opt = Adam(self.model_gen.parameters(), lr=0.01)
+    disriminator_opt = Adam(self.model_disc.parameters(), lr=0.02)
+    return generator_opt, disriminator_opt
+    
+# example with learning_rate schedulers  
+def configure_optimizers(self):
+    generator_opt = Adam(self.model_gen.parameters(), lr=0.01)
+    disriminator_opt = Adam(self.model_disc.parameters(), lr=0.02)
+    discriminator_sched = CosineAnnealing(discriminator_opt, T_max=10)
+    return [generator_opt, disriminator_opt], [discriminator_sched]
+```  
+
+If you need to control how often those optimizers step or override the default .step() schedule, override 
+the [optimizer_step](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step) hook.   
+
+## Optional Methods    
 
 ### validation_step
 
 ``` {.python}
-def validation_step(self, data_batch, batch_nb)
-```
+# if you have one val dataloader:
+def validation_step(self, batch, batch_nb)   
 
-In this step you'd normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something specific to your model.
-This is most likely the same as your training_step. But unlike training step, the outputs from here will go to validation_end for collation.
+# if you have multiple val dataloaders:  
+def validation_step(self, batch, batch_nb, dataloader_idxdx)
+```
+**OPTIONAL**    
+If you don't need to validate you don't need to implement this method. In this step you'd normally generate examples or calculate anything of interest such as accuracy. 
+
+When the validation_step is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, model goes back to training mode and gradients are enabled.
+
+The dict you return here will be available in the `validation_end` method. 
 
 **Params**   
 
 | Param  | description  |
 |---|---|
-|  data_batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch | The output of your dataloader. A tensor, tuple or list  |
 |  batch_nb | Integer displaying which batch this is  |
+|  dataloader_idx | Integer displaying which dataloader this is (only if multiple val datasets used)  |
 
 **Return**   
 
 | Return  | description  | optional |
 |---|---|---|   
-|  dict | Dict of OrderedDict with metrics to display in progress bar. All keys must be tensors. | Y |
+|  dict | Dict or OrderedDict - passed to the validation_end step | N |
 
 **Example**
 
 ``` {.python}
-def validation_step(self, data_batch, batch_nb):
-    x, y, z = data_batch
+# CASE 1: A single validation dataset
+def validation_step(self, batch, batch_nb):
+    x, y = batch
     
     # implement your own
     out = self.forward(x)
-    loss = self.loss(out, x)
+    loss = self.loss(out, y)
+    
+    # log 6 example images
+    # or generated text... or whatever
+    sample_imgs = x[:6]
+    grid = torchvision.utils.make_grid(sample_imgs)
+    self.logger.experiment.add_image('example_images', grid, 0) 
     
     # calculate acc
     labels_hat = torch.argmax(out, dim=1)
@@ -183,31 +396,47 @@ def validation_step(self, data_batch, batch_nb):
     
     # return an optional dict
     return output
-```
+```   
+
+If you pass in multiple validation datasets, validation_step will have an additional argument.
+
+```python
+# CASE 2: multiple validation datasets
+def validation_step(self, batch, batch_nb, dataset_idx):
+    # dataset_idx tells you which dataset this is.   
+```   
+
+The ```dataset_idx``` corresponds to the order of datasets returned in ```val_dataloader```.    
 
 --- 
 ### validation_end
 
 ``` {.python}
 def validation_end(self, outputs)
-```
-
-Called at the end of the validation loop with the output of each validation_step.
+```   
+If you didn't define a validation_step, this won't be called. Called at the end of the validation loop with the outputs of validation_step.
 
+The outputs here are strictly for the progress bar. If you don't need to display anything, don't return anything.    
+Any keys present in 'log', 'progress_bar' or the rest of the dictionary are available for callbacks to access.
 **Params**    
 
 | Param  | description  |
 |---|---|
-|  outputs | List of outputs you defined in validation_step |
+|  outputs | List of outputs you defined in validation_step, or if there are multiple dataloaders, a list containing a list of outputs for each dataloader |
 
 **Return**   
 
-| Return  | description  | optional |
-|---|---|---|   
-|  dict | Dict of OrderedDict with metrics to display in progress bar | Y |
+Dictionary or OrderedDict   
+
+| key  | value  | is required |
+|---|---|---|
+|  progress_bar | Dict for progress bar display. Must have only tensors  | N |
+|  log | Dict of metrics to add to logger. Must have only tensors (no images, etc)  | N |
 
 **Example**
 
+With a single dataloader
+
 ``` {.python}
 def validation_end(self, outputs):
     """
@@ -223,38 +452,197 @@ def validation_end(self, outputs):
 
     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
+    tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+       
+    # show val_loss and val_acc in progress bar but only log val_loss
+    results = {
+        'progress_bar': tqdm_dict,
+        'log': {'val_loss': val_loss_mean.item()}
+    }
+    return results
 ```
 
---- 
-### configure_optimizers 
+With multiple dataloaders, `outputs` will be a list of lists. The outer list contains
+one entry per dataloader, while the inner list contains the individual outputs of 
+each validation step for that dataloader.
 
 ``` {.python}
-def configure_optimizers(self)
+def validation_end(self, outputs):
+    """
+    Called at the end of validation to aggregate outputs
+    :param outputs: list of list of individual outputs of each validation step
+    :return:
+    """
+    val_loss_mean = 0
+    val_acc_mean = 0
+    i = 0
+    for dataloader_outputs in outputs:
+        for output in dataloader_outputs:
+            val_loss_mean += output['val_loss']
+            val_acc_mean += output['val_acc']
+            i += 1
+
+    val_loss_mean /= i
+    val_acc_mean /= i
+    tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+    
+    # show val_loss and val_acc in progress bar but only log val_loss
+    results = {
+        'progress_bar': tqdm_dict,
+        'log': {'val_loss': val_loss_mean.item()}
+    }
+    return results
 ```
 
-Set up as many optimizers and (optionally) learning rate schedulers as you need. Normally you'd need one. But in the case of GANs or something more esoteric you might have multiple. 
-Lightning will call .backward() and .step() on each one in every epoch.  If you use 16 bit precision it will also handle that.
+### test_step
 
+``` {.python}
+# if you have one test dataloader:
+def test_step(self, batch, batch_nb)   
 
-##### Return
-List or Tuple - List of optimizers with an optional second list of learning-rate schedulers
+# if you have multiple test dataloaders:  
+def test_step(self, batch, batch_nb, dataloader_idxdx)
+```
+**OPTIONAL**    
+If you don't need to test you don't need to implement this method. In this step you'd normally generate examples or calculate anything of interest such as accuracy.   
+
+When the validation_step is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, model goes back to training mode and gradients are enabled.
+
+The dict you return here will be available in the `test_end` method. 
+
+This function is used when you execute `trainer.test()`.
+
+**Params**   
+
+| Param  | description  |
+|---|---|
+|  batch | The output of your dataloader. A tensor, tuple or list  |
+|  batch_nb | Integer displaying which batch this is  |
+|  dataloader_idx | Integer displaying which dataloader this is (only if multiple test datasets used)  |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict or OrderedDict with metrics to display in progress bar. All keys must be tensors. | Y |
 
 **Example**
 
 ``` {.python}
-# most cases
-def configure_optimizers(self):
-    opt = Adam(self.parameters(), lr=0.01)
-    return [opt]
+# CASE 1: A single test dataset
+def test_step(self, batch, batch_nb):
+    x, y = batch
     
-# gan example, with scheduler for discriminator
-def configure_optimizers(self):
-    generator_opt = Adam(self.model_gen.parameters(), lr=0.01)
-    disriminator_opt = Adam(self.model_disc.parameters(), lr=0.02)
-    discriminator_sched = CosineAnnealing(discriminator_opt, T_max=10)
-    return [generator_opt, disriminator_opt], [discriminator_sched]
+    # implement your own
+    out = self.forward(x)
+    loss = self.loss(out, y)
+    
+    # calculate acc
+    labels_hat = torch.argmax(out, dim=1)
+    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    
+    # all optional...
+    # return whatever you need for the collation function test_end
+    output = OrderedDict({
+        'test_loss': loss_test,
+        'test_acc': torch.tensor(test_acc), # everything must be a tensor
+    })
+    
+    # return an optional dict
+    return output
+```   
+
+If you pass in multiple test datasets, test_step will have an additional argument.
+
+```python
+# CASE 2: multiple test datasets
+def test_step(self, batch, batch_nb, dataset_idx):
+    # dataset_idx tells you which dataset this is.   
+```   
+
+The ```dataset_idx``` corresponds to the order of datasets returned in ```test_dataloader```.    
+
+--- 
+### test_end
+
+``` {.python}
+def test_end(self, outputs)
+```   
+If you didn't define a test_step, this won't be called.       
+
+Called at the end of the test step with the output of each test_step.
+
+The outputs here are strictly for the progress bar. If you don't need to display anything, don't return anything.    
+
+**Params**    
+
+| Param  | description  |
+|---|---|
+|  outputs | List of outputs you defined in test_step, or if there are multiple dataloaders, a list containing a list of outputs for each dataloader |
+
+**Return**   
+
+| Return  | description  | optional |
+|---|---|---|   
+|  dict | Dict of OrderedDict with metrics to display in progress bar | Y |
+
+**Example**
+
+``` {.python}
+def test_end(self, outputs):
+    """
+    Called at the end of test to aggregate outputs
+    :param outputs: list of individual outputs of each test step
+    :return:
+    """
+    test_loss_mean = 0
+    test_acc_mean = 0
+    for output in outputs:
+        test_loss_mean += output['test_loss']
+        test_acc_mean += output['test_acc']
+
+    test_loss_mean /= len(outputs)
+    test_acc_mean /= len(outputs)
+    tqdm_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+    
+    # show test_loss and test_acc in progress bar but only log test_loss
+    results = {
+        'progress_bar': tqdm_dict,
+        'log': {'test_loss': val_loss_mean.item()}
+    }
+    return results
+```
+
+With multiple dataloaders, `outputs` will be a list of lists. The outer list contains
+one entry per dataloader, while the inner list contains the individual outputs of 
+each validation step for that dataloader.
+
+``` {.python}
+def test_end(self, outputs):
+    """
+    Called at the end of test to aggregate outputs
+    :param outputs: list of individual outputs of each test step
+    :return:
+    """
+    test_loss_mean = 0
+    test_acc_mean = 0
+    i = 0
+    for dataloader_outputs in outputs:
+        for output in dataloader_outputs:
+            test_loss_mean += output['test_loss']
+            test_acc_mean += output['test_acc']
+            i += 1
+
+    test_loss_mean /= i 
+    test_acc_mean /= i
+    tqdm_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+    
+    # show test_loss and test_acc in progress bar but only log test_loss
+    results = {
+        'progress_bar': tqdm_dict,
+        'log': {'test_loss': val_loss_mean.item()}
+    }
+    return results
 ```
 
 --- 
@@ -299,44 +687,21 @@ def on_load_checkpoint(self, checkpoint):
     self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']
 ```
 
---- 
-### tng_dataloader 
-
-``` {.python}
-@pl.data_loader
-def tng_dataloader(self)
-```
-Called by lightning during training loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.
-
-##### Return
-PyTorch DataLoader
-
-**Example**
-
-``` {.python}
-@pl.data_loader
-def tng_dataloader(self):
-    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-    dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform, download=True)
-    loader = torch.utils.data.DataLoader(
-        dataset=dataset,
-        batch_size=self.hparams.batch_size,
-        shuffle=True
-    )
-    return loader
-```
-
 --- 
 ### val_dataloader 
 
 ``` {.python}
 @pl.data_loader
-def tng_dataloader(self)
+def val_dataloader(self)
 ```
-Called by lightning during validation loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.
+**OPTIONAL**    
+If you don't need a validation dataset and a validation_step, you don't need to implement this method.    
+
+Called by lightning during validation loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.   
+If you want to change the data during every epoch DON'T use the data_loader decorator.   
 
 ##### Return
-PyTorch DataLoader
+PyTorch DataLoader or list of PyTorch Dataloaders.    
 
 **Example**
 
@@ -352,8 +717,16 @@ def val_dataloader(self):
     )
     
     return loader
+
+# can also return multiple dataloaders   
+@pl.data_loader
+def val_dataloader(self):
+    return [loader_a, loader_b, ..., loader_n]   
 ```
 
+In the case where you return multiple val_dataloaders, the validation_step will have an arguement ```dataset_idx```
+which matches the order here.    
+
 --- 
 ### test_dataloader 
 
@@ -361,7 +734,11 @@ def val_dataloader(self):
 @pl.data_loader
 def test_dataloader(self)
 ```
+**OPTIONAL**    
+If you don't need a test dataset and a test_step, you don't need to implement this method.    
+
 Called by lightning during test loop. Make sure to use the @pl.data_loader decorator, this ensures not calling this function until the data are needed.
+If you want to change the data during every epoch DON'T use the data_loader decorator.   
 
 ##### Return
 PyTorch DataLoader
@@ -382,26 +759,6 @@ def test_dataloader(self):
     return loader
 ```
 
---- 
-### update_tng_log_metrics 
-
-``` {.python}
-def update_tng_log_metrics(self, logs)
-```
-Called by lightning right before it logs metrics for this batch.
-This is a chance to ammend or add to the metrics about to be logged.
-
-##### Return
-Dict 
-
-**Example**
-
-``` {.python}
-def update_tng_log_metrics(self, logs):
-    # modify or add to logs
-    return logs
-```
-
 --- 
 ### add_model_specific_args 
 
@@ -424,7 +781,7 @@ 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)
+    # parser.set_defaults(gradient_clip_val=5.0)
 
     # network params
     parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)

docs/LightningModule/methods.md

@@ -10,8 +10,25 @@ model.freeze()
 
 ---    
 ### load_from_metrics
-This is the easiest/fastest way which uses the meta_tags.csv file from test-tube to rebuild the model.
-The meta_tags.csv file can be found in the test-tube experiment save_dir.       
+This is the easiest/fastest way which loads hyperparameters and weights from a checkpoint,
+such as the one saved by the `ModelCheckpoint` callback
+
+```{.python}
+pretrained_model = MyLightningModule.load_from_checkpoint(
+    checkpoint_path='/path/to/pytorch_checkpoint.ckpt'
+)
+    
+# predict
+pretrained_model.eval()
+pretrained_model.freeze()
+y_hat = pretrained_model(x)
+```
+
+---    
+### load_from_metrics
+If you're using test tube, there is an alternate method which uses the meta_tags.csv
+file from test-tube to rebuild the model. The meta_tags.csv file can be found in the
+test-tube experiment save_dir.       
 
 ```{.python}
 pretrained_model = MyLightningModule.load_from_metrics(

docs/LightningModule/properties.md

@@ -9,12 +9,22 @@ The current epoch
 Current dtype    
 
 ---
-#### experiment    
-An instance of test-tube Experiment which you can use to log anything for tensorboarX.   
+#### logger    
+A reference to the logger you passed into trainer. 
+Passing a logger is optional. If you don't pass one in, Lightning will create one for you automatically.
+This logger saves logs to '''/os.getcwd()/lightning_logs'''
+```python
+Trainer(logger=your_logger)
+```    
+
+Call it from anywhere in your LightningModule to add metrics, images, etc... whatever your logger supports.  
+
+Here is an example using the TestTubeLogger (which is a wrapper on [PyTorch SummaryWriter](https://pytorch.org/docs/stable/tensorboard.html) with versioned folder structure).     
 ```{.python}
-self.experiment.add_embedding(...)   
-self.experiment.log({'val_loss': 0.9})   
-self.experiment.add_scalars(...)   
+# if logger is a tensorboard logger or TestTubeLogger
+self.logger.experiment.add_embedding(...)   
+self.logger.experiment.log({'val_loss': 0.9})   
+self.logger.experiment.add_scalars(...)   
 ```
 
 --- 
@@ -22,7 +32,7 @@ self.experiment.add_scalars(...)
 Total training batches seen across all epochs   
 
 --- 
-#### gradient_clip   
+#### gradient_clip_val   
 The current gradient clip value    
 
 ---

docs/Trainer/Checkpointing.md

@@ -2,17 +2,25 @@ Lightning can automate saving and loading checkpoints.
 
 ---
 ### Model saving
-To enable checkpointing, define the checkpoint callback and give it to the trainer.
+Checkpointing is enabled by default to the current working directory.
+To change the checkpoint path pass in :
+```python
+Trainer(default_save_path='/your/path/to/save/checkpoints')
+```
+
+To modify the behavior of checkpointing pass in your own callback.
 
 ``` {.python}
-from pytorch_lightning.utils.pt_callbacks import ModelCheckpoint
+from pytorch_lightning.callbacks import ModelCheckpoint
 
+# DEFAULTS used by the Trainer
 checkpoint_callback = ModelCheckpoint(
-    filepath='/path/to/store/weights.ckpt',
+    filepath=os.getcwd(),
     save_best_only=True,
     verbose=True,
     monitor='val_loss',
-    mode='min'
+    mode='min',
+    prefix=''
 )
 
 trainer = Trainer(checkpoint_callback=checkpoint_callback)
@@ -24,12 +32,19 @@ You might want to not only load a model but also continue training it. Use this
 restore the trainer state as well. This will continue from the epoch and global step you last left off.  
 However, the dataloaders will start from the first batch again (if you shuffled it shouldn't matter).   
 
-Lightning will restore the session if you pass an experiment with the same version and there's a saved checkpoint.   
+Lightning will restore the session if you pass a logger with the same version and there's a saved checkpoint.   
 ``` {.python}
-from test_tube import Experiment
+from pytorch_lightning import Trainer
+from pytorch_lightning.logging import TestTubeLogger
 
-exp = Experiment(version=a_previous_version_with_a_saved_checkpoint)
-trainer = Trainer(experiment=exp)
+logger = TestTubeLogger(
+    save_dir='./savepath',
+    version=1  # An existing version with a saved checkpoint
+)
+trainer = Trainer(
+    logger=logger,
+    default_save_path='./savepath'
+)
 
 # this fit call loads model weights and trainer state
 # the trainer continues seamlessly from where you left off
@@ -38,6 +53,7 @@ trainer.fit(model)
 ```
 
 The trainer restores:  
+
 - global_step    
 - current_epoch    
 - All optimizers    
@@ -65,4 +81,4 @@ for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers):
 
 # uses the model you passed into trainer        
 model.load_state_dict(checkpoint['state_dict'])
-```    
+```    

docs/Trainer/Distributed training.md

@@ -8,13 +8,30 @@ None of the flags below require changing anything about your lightningModel defi
 Lightning supports two backends. DataParallel and DistributedDataParallel. Both can be used for single-node multi-GPU training.
 For multi-node training you must use DistributedDataParallel.   
 
+##### DataParallel (dp)   
+Splits a batch across multiple GPUs on the same node. Cannot be used for multi-node training.   
+
+##### DistributedDataParallel (ddp)   
+Trains a copy of the model on each GPU and only syncs gradients. If used with DistributedSampler, each GPU trains
+on a subset of the full dataset.  
+
+##### DistributedDataParallel-2 (ddp2)   
+Works like DDP, except each node trains a single copy of the model using ALL GPUs on that node. 
+Very useful when dealing with negative samples, etc...
+
 You can toggle between each mode by setting this flag.
 ``` {.python}
-# DEFAULT uses DataParallel
+# DEFAULT (when using single GPU or no GPUs)
+trainer = Trainer(distributed_backend=None)
+
+# Change to DataParallel (gpus > 1)
 trainer = Trainer(distributed_backend='dp')
 
-# change to distributed data parallel
+# change to distributed data parallel (gpus > 1)
 trainer = Trainer(distributed_backend='ddp')
+
+# change to distributed data parallel (gpus > 1)
+trainer = Trainer(distributed_backend='ddp2')
 ```
 
 If you request multiple nodes, the back-end will auto-switch to ddp.
@@ -32,12 +49,31 @@ Below are the possible configurations we support.
 
 | 1 GPU  | 1+ GPUs  | DP  | DDP  | 16-bit  | command |
 |---|---|---|---|---|---|
-| Y  |   |   |   |  | ```Trainer(gpus=[0])``` |
-| Y  |   |   |   | Y | ```Trainer(gpus=[0], use_amp=True)``` |
-|   | Y | Y |   |   | ```Trainer(gpus=[0, ...])``` |
-|   | Y |  | Y  |  | ```Trainer(gpus=[0, ...], distributed_backend='ddp')``` |
-|   | Y |  | Y  | Y | ```Trainer(gpus=[0, ...], distributed_backend='ddp', use_amp=True)``` |
+| Y  |   |   |   |  | ```Trainer(gpus=1)``` |
+| Y  |   |   |   | Y | ```Trainer(gpus=1, use_amp=True)``` |
+|   | Y | Y |   |   | ```Trainer(gpus=k, distributed_backend='dp')``` |
+|   | Y |  | Y  |  | ```Trainer(gpus=k, distributed_backend='ddp')``` |
+|   | Y |  | Y  | Y | ```Trainer(gpus=k, distributed_backend='ddp', use_amp=True)``` |
 
+You also have the option of specifying which GPUs to use by passing a list:   
+
+```python
+# DEFAULT (int) specifies how many GPUs to use.
+Trainer(gpus=k)  
+
+# Above is equivalent to 
+Trainer(gpus=list(range(k)))
+
+# You specify which GPUs (don't use if running on cluster)  
+Trainer(gpus=[0, 1])  
+
+# can also be a string
+Trainer(gpus='0, 1')
+
+# can also be -1 or '-1', this uses all available GPUs
+# this is equivalent to list(range(torch.cuda.available_devices()))
+Trainer(gpus=-1)
+```
 
 ---
 #### CUDA flags   
@@ -49,6 +85,9 @@ Lightning sets these for you automatically, there's NO NEED to do this yourself.
 # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 ```
 
+However, when using a cluster, Lightning will NOT set these flags (and you should not either). 
+SLURM will set these for you.   
+
 ---
 #### 16-bit mixed precision
 16 bit precision can cut your memory footprint by half. If using volta architecture GPUs it can give a dramatic training speed-up as well.    
@@ -56,6 +95,21 @@ First, install apex (if install fails, look [here](https://github.com/NVIDIA/ape
 ```bash
 $ git clone https://github.com/NVIDIA/apex
 $ cd apex
+
+# ------------------------
+# OPTIONAL: on your cluster you might need to load cuda 10 or 9
+# depending on how you installed PyTorch
+
+# see available modules
+module avail
+
+# load correct cuda before install
+module load cuda-10.0
+# ------------------------
+
+# make sure you've loaded a cuda version > 4.0 and < 7.0
+module load gcc-6.1.0
+
 $ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
 ```
 
@@ -70,7 +124,7 @@ trainer = Trainer(amp_level='O2', use_amp=False)
 Make sure you're on a GPU machine. 
 ```python
 # DEFAULT
-trainer = Trainer(gpus=[0])
+trainer = Trainer(gpus=1)
 ```
 
 ---
@@ -78,50 +132,63 @@ trainer = Trainer(gpus=[0])
 Make sure you're on a GPU machine. You can set as many GPUs as you want.
 In this setting, the model will run on all 8 GPUs at once using DataParallel under the hood.
 ```python
-# to use DataParallel (default)
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='dp')
+# to use DataParallel
+trainer = Trainer(gpus=8, distributed_backend='dp')
 
 # RECOMMENDED use DistributedDataParallel
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='ddp')
+trainer = Trainer(gpus=8, distributed_backend='ddp')
 ```
 
 ---
 #### Multi-node
-Multi-node training is easily done by specifying these flags.
+Multi-node training is easily done by specifying these flags. 
 ```python
 # train on 12*8 GPUs
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], nb_gpu_nodes=12)
+trainer = Trainer(gpus=8, nb_gpu_nodes=12, distributed_backend='ddp')
 ```
 
-In addition, make sure to set up your SLURM job correctly via the [SlurmClusterObject](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/). In particular, specify the number of tasks per node correctly.
+You must configure your job submission script correctly for the trainer to work. Here is an example
+script for the above trainer configuration.   
 
-```python
-cluster = SlurmCluster(
-    hyperparam_optimizer=test_tube.HyperOptArgumentParser(),
-    log_path='/some/path/to/save',
-)
+```sh
+#!/bin/bash -l
 
-# OPTIONAL FLAGS WHICH MAY BE CLUSTER DEPENDENT
-# which interface your nodes use for communication
-cluster.add_command('export NCCL_SOCKET_IFNAME=^docker0,lo')
+# SLURM SUBMIT SCRIPT
+#SBATCH --nodes=12
+#SBATCH --gres=gpu:8
+#SBATCH --ntasks-per-node=8
+#SBATCH --mem=0
+#SBATCH --time=0-02:00:00
 
-# see output of the NCCL connection process
-# NCCL is how the nodes talk to each other
-cluster.add_command('export NCCL_DEBUG=INFO')
+# activate conda env
+conda activate my_env
 
-# setting a master port here is a good idea.
-cluster.add_command('export MASTER_PORT=%r' % PORT)
+# -------------------------
+# OPTIONAL
+# -------------------------
+# debugging flags (optional)
+# export NCCL_DEBUG=INFO
+# export PYTHONFAULTHANDLER=1
 
-# good to load the latest NCCL version
-cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0'])
+# PyTorch comes with prebuilt NCCL support... but if you have issues with it
+# you might need to load the latest version from your  modules
+# module load NCCL/2.4.7-1-cuda.10.0
 
-# configure cluster
-cluster.per_experiment_nb_nodes = 12 
-cluster.per_experiment_nb_gpus = 8
+# on your cluster you might need these:
+# set the network interface
+# export NCCL_SOCKET_IFNAME=^docker0,lo
+# -------------------------
 
-cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')
+# random port between 12k and 20k
+export MASTER_PORT=$((12000 + RANDOM % 20000))
+
+# run script from above
+python my_main_file.py
 ```
 
+**NOTE:** When running in DDP mode, any errors in your code will show up as an NCCL issue.
+Set the ```NCCL_DEBUG=INFO``` flag to see the ACTUAL error.
+
 Finally, make sure to add a distributed sampler to your dataset. The distributed sampler copies a 
 portion of your dataset onto each GPU. (World_size = gpus_per_node * nb_nodes).   
 
@@ -136,6 +203,58 @@ dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
 dataloader = Dataloader(dataset, sampler=dist_sampler)
 ```
 
+#### Auto-slurm-job-submission
+Instead of manually building SLURM scripts, you can use the [SlurmCluster object](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/) to
+do this for you. The SlurmCluster can also run a grid search if you pass in a [HyperOptArgumentParser](https://williamfalcon.github.io/test-tube/hyperparameter_optimization/HyperOptArgumentParser/).
+
+Here is an example where you run a grid search of 9 combinations of hyperparams.
+[The full examples are here](https://github.com/williamFalcon/pytorch-lightning/tree/master/pl_examples/new_project_templates/multi_node_examples).
+```python
+# grid search 3 values of learning rate and 3 values of number of layers for your net
+# this generates 9 experiments (lr=1e-3, layers=16), (lr=1e-3, layers=32), (lr=1e-3, layers=64), ... (lr=1e-1, layers=64)
+parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+parser.opt_list('--learning_rate', default=0.001, type=float, options=[1e-3, 1e-2, 1e-1], tunable=True)
+parser.opt_list('--layers', default=1, type=float, options=[16, 32, 64], tunable=True)
+hyperparams = parser.parse_args()
+
+# Slurm cluster submits 9 jobs, each with a set of hyperparams
+cluster = SlurmCluster(
+    hyperparam_optimizer=hyperparams,
+    log_path='/some/path/to/save',
+)
+
+# OPTIONAL FLAGS WHICH MAY BE CLUSTER DEPENDENT
+# which interface your nodes use for communication
+cluster.add_command('export NCCL_SOCKET_IFNAME=^docker0,lo')
+
+# see output of the NCCL connection process
+# NCCL is how the nodes talk to each other
+cluster.add_command('export NCCL_DEBUG=INFO')
+
+# setting a master port here is a good idea.
+cluster.add_command('export MASTER_PORT=%r' % PORT)
+
+# ************** DON'T FORGET THIS ***************
+# MUST load the latest NCCL version
+cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0'])
+
+# configure cluster
+cluster.per_experiment_nb_nodes = 12 
+cluster.per_experiment_nb_gpus = 8
+
+cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')  
+
+# submit a script with 9 combinations of hyper params
+# (lr=1e-3, layers=16), (lr=1e-3, layers=32), (lr=1e-3, layers=64), ... (lr=1e-1, layers=64)
+cluster.optimize_parallel_cluster_gpu(
+    main,
+    nb_trials=9, # how many permutations of the grid search to run
+    job_name='name_for_squeue'
+)
+```
+
+The other option is that you generate scripts on your own via a bash command or use another library...
+
 ---
 #### Self-balancing architecture
 Here lightning distributes parts of your module across available GPUs to optimize for speed and memory.   

docs/Trainer/Logging.md

@@ -1,11 +1,164 @@
-Lighting offers a few options for logging information about model, gpu usage, etc (via test-tube). It also offers printing options for training monitoring.
+Lighting offers options for logging information about model, gpu usage, etc, via several different logging frameworks. It also offers printing options for training monitoring.
+
+---   
+### default_save_path   
+Lightning sets a default TestTubeLogger and CheckpointCallback for you which log to
+```os.getcwd()``` by default. To modify the logging path you can set:
+```python
+Trainer(default_save_path='/your/path/to/save/checkpoints')
+```
+ 
+If you need more custom behavior (different paths for both, different metrics, etc...)
+from the logger and the checkpointCallback, pass in your own instances as explained below.
 
 
 ---
+### Setting up logging
+
+The trainer inits a default logger for you (TestTubeLogger). All logs will
+go to the current working directory under a folder named ```os.getcwd()/lightning_logs``. 
+
+If you want to modify the default logging behavior even more, pass in a logger
+(which should inherit from `LightningBaseLogger`).   
+
+```{.python}
+my_logger = MyLightningLogger(...)
+trainer = Trainer(logger=my_logger)
+```
+
+The path in this logger will overwrite default_save_path.
+
+Lightning supports several common experiment tracking frameworks out of the box
+
+---
+#### Test tube
+
+Log using [test tube](https://williamfalcon.github.io/test-tube/). Test tube logger is
+a strict subclass of [PyTorch SummaryWriter](https://pytorch.org/docs/stable/tensorboard.html), refer to their
+documentation for all supported operations. The TestTubeLogger adds a nicer folder structure
+to manage experiments and snapshots all hyperparameters you pass to a LightningModule.
+
+```{.python}
+from pytorch_lightning.logging import TestTubeLogger
+tt_logger = TestTubeLogger(
+    save_dir=".",
+    name="default",
+    debug=False,
+    create_git_tag=False
+)
+trainer = Trainer(logger=tt_logger)
+```
+
+Use the logger anywhere in you LightningModule as follows:
+```python
+def train_step(...):
+    # example
+    self.logger.experiment.whatever_method_summary_writer_supports(...)
+    
+def any_lightning_module_function_or_hook(...):
+    self.logger.experiment.add_histogram(...)
+```
+
+---
+#### MLFlow
+
+Log using [mlflow](https://mlflow.org)
+
+```{.python}
+from pytorch_lightning.logging import MLFlowLogger
+mlf_logger = MLFlowLogger(
+    experiment_name="default",
+    tracking_uri="file:/."
+)
+trainer = Trainer(logger=mlf_logger)
+```
+Use the logger anywhere in you LightningModule as follows:
+```python
+def train_step(...):
+    # example
+    self.logger.experiment.whatever_ml_flow_supports(...)
+    
+def any_lightning_module_function_or_hook(...):
+    self.logger.experiment.whatever_ml_flow_supports(...)
+```
+
+---
+#### Comet.ml
+
+Log using [comet](https://www.comet.ml)
+
+```{.python}
+from pytorch_lightning.logging import CometLogger
+# arguments made to CometLogger are passed on to the comet_ml.Experiment class
+comet_logger = CometLogger(
+    api_key=os.environ["COMET_KEY"],
+    workspace=os.environ["COMET_KEY"],
+)
+trainer = Trainer(logger=comet_logger)
+```
+Use the logger anywhere in you LightningModule as follows:
+```python
+def train_step(...):
+    # example
+    self.logger.experiment.whatever_comet_ml_supports(...)
+
+def any_lightning_module_function_or_hook(...):
+    self.logger.experiment.whatever_comet_ml_supports(...)
+```
+
+---
+#### Custom logger
+
+You can implement your own logger by writing a class that inherits from
+`LightningLoggerBase`. Use the `rank_zero_only` decorator to make sure that
+only the first process in DDP training logs data.
+
+```{.python}
+from pytorch_lightning.logging import LightningLoggerBase, rank_zero_only
+
+class MyLogger(LightningLoggerBase):
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        # params is an argparse.Namespace
+        # your code to record hyperparameters goes here
+        pass
+    
+    @rank_zero_only
+    def log_metrics(self, metrics, step_num):
+        # metrics is a dictionary of metric names and values
+        # your code to record metrics goes here
+        pass
+    
+    def save(self):
+        # Optional. Any code necessary to save logger data goes here
+        pass
+    
+    @rank_zero_only
+    def finalize(self, status):
+        # Optional. Any code that needs to be run after training
+        # finishes goes here
+```
+
+If you write a logger than may be useful to others, please send
+a pull request to add it to Lighting!
+
+---
+#### Using loggers
+You can call the logger anywhere from your LightningModule by doing:
+```python
+def train_step(...):
+    # example
+    self.logger.experiment.whatever_method_summary_writer_supports(...)
+    
+def any_lightning_module_function_or_hook(...):
+    self.logger.experiment.add_histogram(...)
+```
+
 #### Display metrics in progress bar 
 ``` {.python}
 # DEFAULT
-trainer = Trainer(progress_bar=True)
+trainer = Trainer(show_progress_bar=True)
 ```
 
 ---
@@ -13,7 +166,21 @@ trainer = Trainer(progress_bar=True)
 Every k batches lightning will make an entry in the metrics log
 ``` {.python}
 # DEFAULT (ie: save a .csv log file every 10 batches)
-trainer = Trainer(add_log_row_interval=10)
+trainer = Trainer(row_log_interval=10)
+```   
+
+---
+#### Log GPU memory
+Logs GPU memory when metrics are logged.   
+``` {.python}
+# DEFAULT
+trainer = Trainer(log_gpu_memory=None)
+
+# log only the min/max utilization
+trainer = Trainer(log_gpu_memory='min_max')
+
+# log all the GPU memory (if on DDP, logs only that node)
+trainer = Trainer(log_gpu_memory='all')
 ```
 
 ---
@@ -30,61 +197,21 @@ trainer = Trainer(process_position=1)
 
 ---
 #### Save a snapshot of all hyperparameters 
-Whenever you call .save() on the test-tube experiment it logs all the hyperparameters in current use.
-Give lightning a test-tube Experiment object to automate this for you.
+Automatically log hyperparameters stored in the `hparams` attribute as an `argparse.Namespace` 
 ``` {.python}
-from test_tube import Experiment
 
-exp = Experiment(...)
-Trainer(experiment=exp)
-```
+class MyModel(pl.Lightning):
+    def __init__(self, hparams):
+        self.hparams = hparams
 
----
-#### Snapshot code for a training run
-Whenever you call .save() on the test-tube experiment it snapshows all code and pushes to a git tag.
-Give lightning a test-tube Experiment object to automate this for you.
-``` {.python}
-from test_tube import Experiment
+    ...
 
-exp = Experiment(create_git_tag=True)
-Trainer(experiment=exp)
-```
+args = parser.parse_args()
+model = MyModel(args)
 
----
-### Tensorboard support   
-In the LightningModule you can access the experiment logger by doing:
-```python
-self.experiment
-
-# add image
-# Look at PyTorch SummaryWriter docs for what you can do.   
-self.experiment.add_image(...)
-```
-
-The experiment object is a strict subclass of PyTorch SummaryWriter. However, this class
-also snapshots every detail about the experiment (data folder paths, code, hyperparams),
-and allows you to visualize it using tensorboard.
-``` {.python}
-from test_tube import Experiment, HyperOptArgumentParser
-
-# exp hyperparams
-args = HyperOptArgumentParser()
-hparams = args.parse_args()
-
-# this is a summaryWriter with nicer logging structure
-exp = Experiment(save_dir='/some/path', create_git_tag=True)
-
-# track experiment details (must be ArgumentParser or HyperOptArgumentParser).
-# each option in the parser is tracked
-exp.argparse(hparams)
-exp.tag({'description': 'running demo'})
-
-# trainer uses the exp object to log exp data
-trainer = Trainer(experiment=exp)
+logger = TestTubeLogger(...)
+t = Trainer(logger=logger)
 trainer.fit(model)
-
-# view logs at:
-# tensorboard --logdir /some/path   
 ```
 
 ---

docs/Trainer/SLURM Managed Cluster.md

@@ -1,8 +1,10 @@
 Lightning supports model training on a cluster managed by SLURM in the following cases:    
 
-1. Training on single or multi-cpus only.
-2. Training on single or multi-gpus on the same node.
-3. Coming SOON: Training across multiple nodes.
+1. Training on a single cpu or single GPU.
+2. Train on multiple GPUs on the same node using DataParallel or DistributedDataParallel
+3. Training across multiple GPUs on multiple different nodes via DistributedDataParallel.
+
+**Note: A node means a machine with multiple GPUs**
 
 ---
 #### Running grid search on a cluster
@@ -23,6 +25,9 @@ parser.opt_list('--nb_layers', default=2, type=int, tunable=True, options=[2, 4,
 hparams = parser.parse_args()    
 ```    
     
+**NOTE** You must set ```Tunable=True``` for that argument to be considered in the permutation set. Otherwise
+test-tube will use the default value. This flag is useful when you don't want to search over an argument and
+want to use the default instead.   
      
 (2). Define the cluster options in the [SlurmCluster object](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/) (over 5 nodes and 8 gpus)    
 
@@ -55,8 +60,8 @@ cluster.memory_mb_per_node = 10000
 cluster.job_time = '10:00'
 ```
 
-(3). Give trainer the cluster_manager in your main function:    
-
+(3). Make a main function with your model and trainer. Each job will call this function with a particular
+hparams configuration.    
 ```{.python}
 from pytorch_lightning import Trainer
 
@@ -66,12 +71,12 @@ def train_fx(trial_hparams, cluster_manager, _):
     my_model = MyLightningModel()
     
     # give the trainer the cluster object
-    trainer = Trainer(cluster=cluster_manager)
+    trainer = Trainer()
     trainer.fit(my_model)
 
 ```
 
-(4). Start the grid search     
+(3). Start the grid/random search     
 ```{.python}
 # run the models on the cluster
 cluster.optimize_parallel_cluster_gpu(
@@ -81,24 +86,27 @@ cluster.optimize_parallel_cluster_gpu(
     job_display_name='my_exp')
 ```
 
-That's it! The SlurmCluster object will automatically checkpoint the lightning model and resubmit if it runs into the walltime!
-
+**NOTE** nb_trials specifies how many of the possible permutations to use. If using ```grid_search``` it will use
+the depth first ordering. If using ```random_search``` it will use the first k shuffled options. FYI, random search
+has been shown to be just as good as any Bayesian optimization method when using a reasonable number of samples (60),
+[see this paper for more information](http://www.jmlr.org/papers/volume13/bergstra12a/bergstra12a.pdf).
 
 ---
 #### Walltime auto-resubmit
-Lightning automatically resubmits jobs when they reach the walltime. You get this behavior for free if you give lightning
-a slurm cluster object.
+Lightning automatically resubmits jobs when they reach the walltime. Make sure to set the SIGUSR1 signal in 
+your SLURM script.   
 
-```{.python}
-def my_main_fx(hparams, slurm_manager, _):
-    trainer = Trainer(cluster=slurm_manager)
+```bash
+# 90 seconds before training ends
+#SBATCH --signal=SIGUSR1@90
 ``` 
 
-(See the grid search example above for cluster configuration).
-With this feature lightning will:    
+When lightning receives the SIGUSR1 signal it will:
+1. save a checkpoint with 'hpc_ckpt' in the name.
+2. resubmit the job using the SLURM_JOB_ID  
+
+When the script starts again, Lightning will:
+1. search for a 'hpc_ckpt' checkpoint. 
+2. restore the model, optimizers, schedulers, epoch, etc...   
 
-1. automatically checkpoint the model
-2. checkpoint the trainer session
-3. resubmit a continuation job.
-4. load the checkpoint and trainer session in the new model
 

docs/Trainer/Testing loop.md

@@ -0,0 +1,31 @@
+To ensure you don't accidentally use test data to guide training decisions Lightning makes running the test set deliberate.   
+
+---
+#### test
+You have two options to run the test set. 
+First case is where you test right after a full training routine.
+``` {.python}
+# run full training
+trainer.fit(model)
+
+# run test set
+trainer.test()
+```   
+
+Second case is where you load a model and run the test set   
+```{.python}
+model = MyLightningModule.load_from_metrics(
+    weights_path='/path/to/pytorch_checkpoint.ckpt',
+    tags_csv='/path/to/test_tube/experiment/version/meta_tags.csv',
+    on_gpu=True,
+    map_location=None
+)
+
+# init trainer with whatever options
+trainer = Trainer(...)
+    
+# test (pass in the model)
+trainer.test(model)
+```
+In this second case, the options you pass to trainer will be used when running the test set (ie: 16-bit, dp, ddp, etc...)  
+ 

docs/Trainer/Training Loop.md

@@ -3,8 +3,8 @@ The lightning training loop handles everything except the actual computations of
 Below are all the things lightning automates for you in the training loop.
 
 ---
-#### Accumulated gradients  
-Accumulated gradients runs K small batches of size N before doing a backwards pass. The effect is a large effective batch size of size KxN. 
+#### Accumulated gradients
+Accumulated gradients runs K small batches of size N before doing a backwards pass. The effect is a large effective batch size of size KxN.
 
 ``` {.python}
 # DEFAULT (ie: no accumulated grads)
@@ -20,22 +20,51 @@ trainer = Trainer(min_nb_epochs=1, max_nb_epochs=1000)
 ```
 
 ---
-#### Force disable early stop 
-Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
+#### Early stopping
+The trainer already sets up default early stopping for you.
+To modify this behavior, pass in your own EarlyStopping callback.
 ``` {.python}
-# DEFAULT
-trainer = Trainer(enable_early_stop=True)
+from pytorch_lightning.callbacks import EarlyStopping
+
+# DEFAULTS used by Trainer
+early_stop_callback = EarlyStopping(
+    monitor='val_loss',
+    min_delta=0.00,
+    patience=3,
+    verbose=False,
+    mode='min'
+)
+
+# without passing anything in, uses the default callback above
+trainer = Trainer()
+
+# pass in your own to override the default callback
+trainer = Trainer(early_stop_callback=early_stop_callback)
+
+# pass in None to disable it
+trainer = Trainer(early_stop_callback=None)
 ```
 
 ---
-#### Gradient Clipping 
-Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
+#### Force disable early stop
+To disable early stopping pass None to the early_stop_callback
 ``` {.python}
-# DEFAULT (ie: don't clip)
-trainer = Trainer(gradient_clip=0)
+# DEFAULT
+trainer = Trainer(early_stop_callback=None)
 ```
 
+---
+#### Gradient Clipping
+Gradient clipping may be enabled to avoid exploding gradients.
+Specifically, this will [clip the gradient norm computed over all model parameters *together*](https://pytorch.org/docs/stable/nn.html#torch.nn.utils.clip_grad_norm_).
 
+``` {.python}
+# DEFAULT (ie: don't clip)
+trainer = Trainer(gradient_clip_val=0)
+
+# clip gradients with norm above 0.5
+trainer = Trainer(gradient_clip_val=0.5)
+```
 
 ---
 #### Inspect gradient norms
@@ -51,7 +80,10 @@ trainer = Trainer(track_grad_norm=2)
 
 ---
 #### Set how much of the training set to check
-If you don't want to check 100% of the training set (for debugging or if it's huge), set this flag
+If you don't want to check 100% of the training set (for debugging or if it's huge), set this flag.
+
+train_percent_check will be overwritten by overfit_pct if `overfit_pct > 0`
+
 ``` {.python}
 # DEFAULT
 trainer = Trainer(train_percent_check=1.0)
@@ -59,3 +91,36 @@ trainer = Trainer(train_percent_check=1.0)
 # check 10% only
 trainer = Trainer(train_percent_check=0.1)
 ```
+
+---
+#### Packed sequences as inputs
+When using PackedSequence, do 2 things:
+1. return either a padded tensor in dataset or a list of variable length tensors in the dataloader collate_fn (example above shows the list implementation).    
+2. Pack the sequence in forward or training and validation steps depending on use case.
+
+``` {.python}
+# For use in dataloader
+def collate_fn(batch):
+    x = [item[0] for item in batch]
+    y = [item[1] for item in batch]
+    return x, y
+
+# In module
+def training_step(self, batch, batch_nb):
+    x = rnn.pack_sequence(batch[0], enforce_sorted=False)
+    y = rnn.pack_sequence(batch[1], enforce_sorted=False)
+```
+
+---
+#### Truncated Back Propagation Through Time
+There are times when multiple backwards passes are needed for each batch. For example, it may save memory to use Truncated Back Propagation Through Time when training RNNs.
+
+When this flag is enabled each batch is split into sequences of size truncated_bptt_steps and passed to training_step(...) separately. A default splitting function is provided, however, you can override it for more flexibility. See [tbptt_split_batch](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks#tbptt_split_batch).
+
+``` {.python}
+# DEFAULT (single backwards pass per batch)
+trainer = Trainer(truncated_bptt_steps=None)
+
+# (split batch into sequences of size 2)
+trainer = Trainer(truncated_bptt_steps=2)
+```

docs/Trainer/Validation loop.md

@@ -5,8 +5,6 @@ Below are all the things lightning automates for you in the validation loop.
 Lightning will run 5 steps of validation in the beginning of training as a sanity check so you don't have to wait until a full epoch to catch possible validation issues.
 
 
-
-
 ---
 #### Check validation every n epochs
 If you have a small dataset you might want to check validation every n epochs
@@ -18,6 +16,9 @@ trainer = Trainer(check_val_every_n_epoch=1)
 ---
 #### Set how much of the validation set to check 
 If you don't want to check 100% of the validation set (for debugging or if it's huge), set this flag
+
+val_percent_check will be overwritten by overfit_pct if `overfit_pct > 0`
+
 ``` {.python}
 # DEFAULT
 trainer = Trainer(val_percent_check=1.0)
@@ -29,6 +30,9 @@ trainer = Trainer(val_percent_check=0.1)
 ---
 #### Set how much of the test set to check 
 If you don't want to check 100% of the test set (for debugging or if it's huge), set this flag
+
+test_percent_check will be overwritten by overfit_pct if `overfit_pct > 0`
+
 ``` {.python}
 # DEFAULT
 trainer = Trainer(test_percent_check=1.0)
@@ -39,13 +43,20 @@ trainer = Trainer(test_percent_check=0.1)
 
 ---
 ####  Set validation check frequency within 1 training epoch
-For large datasets it's often desirable to check validation multiple times within a training loop
+For large datasets it's often desirable to check validation multiple times within a training loop.
+Pass in a float to check that often within 1 training epoch.
+Pass in an int k to check every k training batches. Must use an int if using 
+an IterableDataset.
+ 
 ``` {.python}
 # DEFAULT
 trainer = Trainer(val_check_interval=0.95)
 
 # check every .25 of an epoch 
 trainer = Trainer(val_check_interval=0.25)
+
+# check every 100 train batches (ie: for IterableDatasets or fixed frequency)
+trainer = Trainer(val_check_interval=100)
 ```
 
 ---
@@ -54,4 +65,6 @@ Lightning runs a few steps of validation in the beginning of training. This avoi
 ``` {.python}
 # DEFAULT
 trainer = Trainer(nb_sanity_val_steps=5)
-```
+```
+
+You can use `Trainer(nb_sanity_val_steps=0)` to skip the sanity check.

docs/Trainer/debugging.md

@@ -23,6 +23,9 @@ trainer = Trainer(track_grad_norm=2)
 ---
 #### Make model overfit on subset of data
 A useful debugging trick is to make your model overfit a tiny fraction of the data.
+
+setting `overfit_pct > 0` will overwrite train_percent_check, val_percent_check, test_percent_check
+
 ``` {.python}
 # DEFAULT don't overfit (ie: normal training)
 trainer = Trainer(overfit_pct=0.0)
@@ -35,6 +38,14 @@ trainer = Trainer(overfit_pct=0.01)
 #### Print the parameter count by layer
 By default lightning prints a list of parameters *and submodules* when it starts training.
 
+``` {.python}
+# DEFAULT print a full list of all submodules and their parameters.
+trainer = Trainer(weights_summary='full')
+
+# only print the top-level modules (i.e. the children of LightningModule).
+trainer = Trainer(weights_summary='top')
+```
+
 ---
 #### Print which gradients are nan 
 This option prints a list of tensors with nan gradients.

docs/Trainer/hooks.md

@@ -58,15 +58,52 @@ def on_post_performance_check(self):
 ```
 
 ---
-#### on_tng_metrics
-Called in the training loop, right before metrics are logged.
-Although you can log at any time by using self.experiment, you can use
-this callback to modify what will be logged.
+#### optimizer_step 
+Calls .step() and .zero_grad for each optimizer.  
+You can override this method to adjust how you do the optimizer step for each optimizer
+
+Called once per optimizer
 ```python
-def on_tng_metrics(self, metrics):
-    # do something before validation end
+# DEFAULT
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i, second_order_closure=None):
+    optimizer.step()   
+    optimizer.zero_grad()   
+    
+# Alternating schedule for optimizer steps (ie: GANs)    
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i, second_order_closure=None):
+    # update generator opt every 2 steps
+    if optimizer_i == 0:
+        if batch_nb % 2 == 0 :
+            optimizer.step()
+            optimizer.zero_grad()
+   
+    # update discriminator opt every 4 steps
+    if optimizer_i == 1:
+        if batch_nb % 4 == 0 :
+            optimizer.step()
+            optimizer.zero_grad()    
+    
+    # ...
+    # add as many optimizers as you want 
 ```
 
+This step allows you to do a lot of non-standard training tricks such as learning-rate warm-up:   
+
+```python
+# learning rate warm-up
+def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i, second_order_closure=None):
+    # warm up lr
+    if self.trainer.global_step < 500:
+        lr_scale = min(1., float(self.trainer.global_step + 1) / 500.)
+        for pg in optimizer.param_groups:
+            pg['lr'] = lr_scale * self.hparams.learning_rate
+    
+    # update params
+    optimizer.step()
+    optimizer.zero_grad() 
+```
+
+
 ---
 #### on_before_zero_grad
 Called in the training loop after taking an optimizer step and before zeroing grads.
@@ -78,6 +115,28 @@ def on_before_zero_grad(self, optimizer):
     # do something with the optimizer or inspect it. 
 ```
 
+---
+#### backward
+Called to perform backward step.
+Feel free to override as needed.
+
+The loss passed in has already been scaled for accumulated gradients if requested.
+```python
+def backward(self, use_amp, loss, optimizer):
+    """
+    Override backward with your own implementation if you need to
+    :param use_amp: Whether amp was requested or not
+    :param loss: Loss is already scaled by accumulated grads
+    :param optimizer: Current optimizer being used
+    :return:
+    """
+    if use_amp:
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+    else:
+        loss.backward()
+```
+
 ---
 #### on_after_backward
 Called in the training loop after model.backward()
@@ -90,5 +149,118 @@ def on_after_backward(self):
         for k, v in params.items():
             grads = v
             name = k
-            self.experiment.add_histogram(tag=name, values=grads, global_step=self.trainer.global_step)
+            self.logger.experiment.add_histogram(tag=name, values=grads, global_step=self.trainer.global_step)
+```
+
+---
+#### tbptt_split_batch
+Called in the training loop after on_batch_start if `truncated_bptt_steps > 0`. Each returned batch split is passed separately to training_step(...).
+
+```python
+def tbptt_split_batch(self, batch, split_size):
+  splits = []
+  for t in range(0, time_dims[0], split_size):
+      batch_split = []
+      for i, x in enumerate(batch):
+          if isinstance(x, torch.Tensor):
+              split_x = x[:, t:t + split_size]
+          elif isinstance(x, collections.Sequence):
+              split_x = [None] * len(x)
+              for batch_idx in range(len(x)):
+                  split_x[batch_idx] = x[batch_idx][t:t + split_size]
+
+          batch_split.append(split_x)
+
+      splits.append(batch_split)
+
+  return splits
+```
+
+---
+#### configure_apex
+Overwrite to define your own Apex implementation init.
+
+```python
+def configure_apex(self, amp, model, optimizers, amp_level):
+    """
+    Override to init AMP your own way
+    Must return a model and list of optimizers
+    :param amp:
+    :param model:
+    :param optimizers:
+    :param amp_level:
+    :return: Apex wrapped model and optimizers
+    """
+    model, optimizers = amp.initialize(
+        model, optimizers, opt_level=amp_level,
+    )
+
+    return model, optimizers
+```
+
+---
+#### configure_ddp 
+Overwrite to define your own DDP implementation init.
+The only requirement is that:
+1. On a validation batch the call goes to model.validation_step.   
+2. On a training batch the call goes to model.training_step.   
+3. On a testing batch, the call goes to model.test_step
+
+```python
+def configure_ddp(self, model, device_ids):
+    """
+    Override to init DDP in a different way or use your own wrapper.
+    Must return model.
+    :param model:
+    :param device_ids:
+    :return: DDP wrapped model
+    """
+    # Lightning DDP simply routes to test_step, val_step, etc...
+    model = LightningDistributedDataParallel(
+        model,
+        device_ids=device_ids,
+        find_unused_parameters=True
+    )
+    return model
+```
+
+---   
+#### init_ddp_connection   
+Override to init DDP in your own way.   
+
+```python
+def init_ddp_connection(self):
+    """
+    Connect all procs in the world using the env:// init
+    Use the first node as the root address
+    """
+
+    # use slurm job id for the port number
+    # guarantees unique ports across jobs from same grid search
+    try:
+        # use the last 4 numbers in the job id as the id
+        default_port = os.environ['SLURM_JOB_ID']
+        default_port = default_port[-4:]
+
+        # all ports should be in the 10k+ range
+        default_port = int(default_port) + 15000
+
+    except Exception as e:
+        default_port = 12910
+
+    # if user gave a port number, use that one instead
+    try:
+        default_port = os.environ['MASTER_PORT']
+    except Exception:
+        os.environ['MASTER_PORT'] = str(default_port)
+
+    # figure out the root node addr
+    try:
+        root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+    except Exception:
+        root_node = '127.0.0.2'
+
+    root_node = self.trainer.resolve_root_node_address(root_node)
+    os.environ['MASTER_ADDR'] = root_node
+    dist.init_process_group('nccl', rank=self.proc_rank, world_size=self.world_size)
 ```

docs/Trainer/index.md

@@ -1,5 +1,5 @@
 # Trainer
-[[Github Code](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/trainer.py)]
+[[Github Code](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/trainer/trainer.py)]
 
 The lightning trainer abstracts best practices for running a training, val, test routine. It calls parts of your model when it wants to hand over full control and otherwise makes training assumptions which are now standard practice in AI research.
 
@@ -19,6 +19,7 @@ But of course the fun is in all the advanced things it can do:
 
 **Checkpointing**    
 
+- [Checkpoint callback](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)    
 - [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
 - [Restoring training session](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#restoring-training-session)
@@ -35,12 +36,13 @@ But of course the fun is in all the advanced things it can do:
 - [Log GPU usage](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#Log-gpu-usage)
 - [Make model overfit on subset of data](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#make-model-overfit-on-subset-of-data)
 - [Print the parameter count by layer](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-the-parameter-count-by-layer)
-- [Pring which gradients are nan](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-which-gradients-are-nan)
+- [Print which gradients are nan](https://williamfalcon.github.io/pytorch-lightning/Trainer/debugging/#print-which-gradients-are-nan)
 - [Print input and output size of every module in system](https://williamfalcon.github.io/pytorch-lightning/LightningModule/properties/#example_input_array)
 
 
 **Distributed training**    
 
+- [Implement Your Own Distributed (DDP) training](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#init_ddp_connection)
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
 - [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
@@ -60,14 +62,18 @@ But of course the fun is in all the advanced things it can do:
 
 **Training loop**    
 
-- [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
+- [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)    
 - [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
+- [Early stopping callback](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#early-stopping)    
 - [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
 - [Gradient Clipping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
 - [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
 - [Learning rate scheduling](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
 - [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
 - [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
+- [Step optimizers at arbitrary intervals](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step)
+- [Packed sequences](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#packed-sequences-as-inputs)
+- [Truncated Back Propagation Through Time](https://williamfalcon.github.io/pytorch-lightning//Training%20Loop/#truncated-back-propation-through-time)
 
 **Validation loop**    
 
@@ -77,3 +83,8 @@ But of course the fun is in all the advanced things it can do:
 - [Set how much of the test set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-test-set-to-check)
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
 - [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
+
+
+**Testing loop**    
+
+- [Run test set](https://williamfalcon.github.io/pytorch-lightning/Trainer/Testing%20loop/)  

docs/examples/Examples.md

@@ -1,9 +1,9 @@
 ### Template model definition
-In 99% of cases you want to just copy [this template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py) to start a new lightningModule and change the core of what your model is actually trying to do.
+In 99% of cases you want to just copy [one of the examples](https://github.com/williamFalcon/pytorch-lightning/tree/master/pl_examples) to start a new lightningModule and change the core of what your model is actually trying to do.
 
 ```bash
 # get a copy of the module template
-wget https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py
+wget https://raw.githubusercontent.com/williamFalcon/pytorch-lightning/master/pl_examples/new_project_templates/lightning_module_template.py
 ```
 
 ---    
@@ -40,64 +40,24 @@ The main function should have 3 arguments:
 - slurm_manager: Slurm cluster manager object (can be None)
 - dict: for you to return any values you want (useful in meta-learning, otherwise set to _)    
 
-```{}
+```python
 def main(hparams, cluster, results_dict):
     """
     Main training routine specific for this project
     :param hparams:
     :return:
     """
-    # 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'
-    )
-    
-    # set the hparams for the experiment
-    exp.argparse(hparams)
-    exp.save()
-
     # build model
     model = MyLightningModule(hparams)
 
-    # 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
-    )
-
     # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        cluster=cluster,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-    )
+    trainer = Trainer()
 
     # train model
     trainer.fit(model)
 ```
 
 
-
-
 The __main__ function will start training on your **main** function. If you use the HyperParameterOptimizer
 in hyper parameter optimization mode, this main function will get one set of hyperparameters. If you use it as a simple
 argument parser you get the default arguments in the argument parser.
@@ -159,7 +119,7 @@ def optimize_on_cluster(hyperparams):
     job_display_name = job_display_name[0:3]
 
     # run hopt
-    print('submitting jobs...')
+    logging.info('submitting jobs...')
     cluster.optimize_parallel_cluster_gpu(
         main,
         nb_trials=hyperparams.nb_hopt_trials,

docs/index.md

@@ -1,17 +1,67 @@
 ###### New project Quick Start    
-To start a new project you define two files, a LightningModule and a Trainer file.    
+To start a new project define two files, a LightningModule and a Trainer file.    
+To illustrate Lightning power and simplicity, here's an example of a typical research flow.    
 
-A separate trainer file allows to run many LightningModules. Each LightningModule has the core
-logic to a particular research project.    
+###### Case 1: BERT    
+Let's say you're working on something like BERT but want to try different ways of training or even different networks.  
+You would define a single LightningModule and use flags to switch between your different ideas.   
+```python
+class BERT(pl.LightningModule):
+    def __init__(self, model_name, task):
+        self.task = task
+    
+        if model_name == 'transformer':
+            self.net = Transformer()
+        elif model_name == 'my_cool_version':
+            self.net = MyCoolVersion()
+            
+    def training_step(self, batch, batch_nb):
+        if self.task == 'standard_bert':
+            # do standard bert training with self.net...
+            # return loss
+            
+        if self.task == 'my_cool_task':
+            # do my own version with self.net
+            # return loss
+```   
 
-For example, one lightningModule could be an image classifier, the other
-one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
+###### Case 2: COOLER NOT BERT    
+But if you wanted to try something **completely** different, you'd define a new module for that.    
+```python
 
+class CoolerNotBERT(pl.LightningModule):
+    def __init__(self):
+        self.net = ...
+        
+    def training_step(self, batch, batch_nb):
+        # do some other cool task
+        # return loss   
+```   
+
+###### Rapid research flow    
+Then you could do rapid research by switching between these two and using the same trainer.   
+```python
+
+if use_bert:
+    model = BERT()
+else:
+    model = CoolerNotBERT()
+    
+trainer = Trainer(gpus=4, use_amp=True)
+trainer.fit(model)
+```
+
+Notice a few things about this flow:   
+1. You're writing pure PyTorch... no unnecessary abstractions or new libraries to learn.   
+2. You get free GPU and 16-bit support without writing any of that code in your model.   
+3. You also get all of the capabilities below (without coding or testing yourself).     
+
+---    
+###### Templates 
 1. [MNIST LightningModule](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#minimal-example) 
 2. [Trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)
-    - [Basic CPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/single_cpu_template.py) 
-    - [Multi-GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/single_gpu_node_template.py)
-    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/multi_node_cluster_template.py)
+    - [Basic CPU, GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/tree/master/pl_examples/basic_examples)
+    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/tree/master/pl_examples/multi_node_examples)
 
 ###### Docs shortcuts
 - [LightningModule](LightningModule/RequiredTrainerInterface/)  
@@ -26,6 +76,7 @@ one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
 
 ###### Checkpointing    
 
+- [Checkpoint callback](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)    
 - [Model saving](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#model-saving)
 - [Model loading](https://williamfalcon.github.io/pytorch-lightning/LightningModule/methods/#load-from-metrics) 
 - [Restoring training session](https://williamfalcon.github.io/pytorch-lightning/Trainer/Checkpointing/#restoring-training-session)
@@ -48,6 +99,7 @@ one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
 
 ###### Distributed training    
 
+- [Implement Your Own Distributed (DDP) training](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#init_ddp_connection)
 - [16-bit mixed precision](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#16-bit-mixed-precision)
 - [Multi-GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-GPU)
 - [Multi-node](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#Multi-node)
@@ -69,12 +121,14 @@ one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
 
 - [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
 - [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
+- [Early stopping callback](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#early-stopping)    
 - [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
 - [Gradient Clipping](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
 - [Hooks](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/)
 - [Learning rate scheduling](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
 - [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#configure_optimizers)
 - [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
+- [Step optimizers at arbitrary intervals](https://williamfalcon.github.io/pytorch-lightning/Trainer/hooks/#optimizer_step)
 
 ###### Validation loop    
 
@@ -85,3 +139,5 @@ one could be a seq-2-seq model, both (optionally) ran by the same trainer file.
 - [Set validation check frequency within 1 training epoch](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-validation-check-frequency-within-1-training-epoch)
 - [Set the number of validation sanity steps](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-the-number-of-validation-sanity-steps)
 
+###### Testing loop  
+- [Run test set](https://williamfalcon.github.io/pytorch-lightning/Trainer/Testing%20loop/)  

examples/__init__.py

@@ -1,5 +0,0 @@
-from .new_project_templates.lightning_module_template import LightningTemplateModel
-
-__all__ = [
-    'LightningTemplateModel'
-]

examples/new_project_templates/multi_node_cluster_template.py

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

examples/new_project_templates/single_cpu_template.py

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

examples/new_project_templates/single_gpu_node_16bit_template.py

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

examples/new_project_templates/single_gpu_node_ddp_template.py

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

examples/new_project_templates/single_gpu_node_dp_template.py

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

examples/new_project_templates/trainer_cpu_template.py

@@ -1,74 +0,0 @@
-import os
-import sys
-
-from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utilities.arg_parse import add_default_args
-from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
-
-from examples.new_project_templates.lightning_module_template import LightningTemplateModel
-
-
-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 = LightningTemplateModel(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_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 = LightningTemplateModel.add_model_specific_args(parent_parser)
-    hyperparams = parser.parse_args()
-
-    # train model
-    main(hyperparams)

mkdocs.yml

@@ -2,9 +2,15 @@ site_name: PyTorch lightning Documentation
 theme:
   name: 'material'
 docs_dir: docs
+repo_name: 'williamFalcon/pytorch-lightning'
 repo_url: https://github.com/williamFalcon/pytorch-lightning
 site_dir: 'site'
 site_description: 'Documentation for PyTorch LightningModule, the researcher version of keras.'
 
 dev_addr: '0.0.0.0:8000'
 #google_analytics: ['UA-aasd', 'sitename']
+
+markdown_extensions:
+  - codehilite:
+      guess_lang: false
+      linenums: true

pl_examples/README.md

@@ -0,0 +1,11 @@
+# Examples   
+This folder has 3 sections:   
+
+### Domain templates   
+These are templates to show common approaches such as GANs and RL.   
+
+### Basic examples   
+These show the most common use of Lightning for either CPU or GPU training.   
+
+### Multi-node examples   
+These show how to run jobs on a GPU cluster using lightning.

pl_examples/__init__.py

@@ -0,0 +1,5 @@
+from .basic_examples.lightning_module_template import LightningTemplateModel
+
+__all__ = [
+    'LightningTemplateModel'
+]

pl_examples/basic_examples/README.md

@@ -0,0 +1,39 @@
+# Basic Examples   
+Use these examples to test how lightning works.   
+
+#### Test on CPU  
+```bash
+python cpu_template.py
+```
+
+---   
+#### Train on a single GPU
+```bash
+python gpu_template.py --gpus 1
+```   
+
+---    
+#### DataParallel (dp)   
+Train on multiple GPUs using DataParallel.
+
+```bash
+python gpu_template.py --gpus 2 --distributed_backend dp
+```   
+
+---
+#### DistributedDataParallel (ddp)    
+
+Train on multiple GPUs using DistributedDataParallel   
+```bash
+python gpu_template.py --gpus 2 --distributed_backend ddp
+```
+
+---
+#### DistributedDataParallel+DP (ddp2)    
+
+Train on multiple GPUs using DistributedDataParallel + dataparallel.
+On a single node, uses all GPUs for 1 model. Then shares gradient information
+across nodes.   
+```bash
+python gpu_template.py --gpus 2 --distributed_backend ddp2
+```

pl_examples/basic_examples/cpu_template.py

@@ -0,0 +1,54 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+from argparse import ArgumentParser
+
+import numpy as np
+import torch
+
+from pl_examples.basic_examples.lightning_module_template import LightningTemplateModel
+from pytorch_lightning import Trainer
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = LightningTemplateModel(hparams)
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = Trainer()
+
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    # ------------------------
+    # TRAINING ARGUMENTS
+    # ------------------------
+    # these are project-wide arguments
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    parent_parser = ArgumentParser(add_help=False)
+
+    # each LightningModule defines arguments relevant to it
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    main(hyperparams)

pl_examples/basic_examples/gpu_template.py

@@ -0,0 +1,79 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+from argparse import ArgumentParser
+
+import numpy as np
+import torch
+
+from pl_examples.basic_examples.lightning_module_template import LightningTemplateModel
+from pytorch_lightning import Trainer
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = LightningTemplateModel(hparams)
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        gpus=hparams.gpus,
+        distributed_backend=hparams.distributed_backend,
+        use_amp=hparams.use_16bit
+    )
+
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    # ------------------------
+    # TRAINING ARGUMENTS
+    # ------------------------
+    # these are project-wide arguments
+
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    parent_parser = ArgumentParser(add_help=False)
+
+    # gpu args
+    parent_parser.add_argument(
+        '--gpus',
+        type=int,
+        default=2,
+        help='how many gpus'
+    )
+    parent_parser.add_argument(
+        '--distributed_backend',
+        type=str,
+        default='dp',
+        help='supports three options dp, ddp, ddp2'
+    )
+    parent_parser.add_argument(
+        '--use_16bit',
+        dest='use_16bit',
+        action='store_true',
+        help='if true uses 16 bit precision'
+    )
+
+    # each LightningModule defines arguments relevant to it
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    main(hyperparams)

pl_examples/basic_examples/lightning_module_template.py

@@ -2,16 +2,18 @@
 Example template for defining a system
 """
 import os
+import logging
+from argparse import ArgumentParser
 from collections import OrderedDict
-import torch.nn as nn
-from torchvision.datasets import MNIST
-import torchvision.transforms as transforms
+
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
-from test_tube import HyperOptArgumentParser
+import torchvision.transforms as transforms
 from torch import optim
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
+from torchvision.datasets import MNIST
 
 import pytorch_lightning as pl
 from pytorch_lightning.root_module.root_module import LightningModule
@@ -79,14 +81,14 @@ class LightningTemplateModel(LightningModule):
         nll = F.nll_loss(logits, labels)
         return nll
 
-    def training_step(self, data_batch, batch_i):
+    def training_step(self, batch, batch_idx):
         """
         Lightning calls this inside the training loop
-        :param data_batch:
+        :param batch:
         :return:
         """
         # forward pass
-        x, y = data_batch
+        x, y = batch
         x = x.view(x.size(0), -1)
 
         y_hat = self.forward(x)
@@ -95,23 +97,26 @@ class LightningTemplateModel(LightningModule):
         loss_val = self.loss(y, y_hat)
 
         # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
+        if self.trainer.use_dp or self.trainer.use_ddp2:
             loss_val = loss_val.unsqueeze(0)
 
+        tqdm_dict = {'train_loss': loss_val}
         output = OrderedDict({
-            'loss': loss_val
+            'loss': loss_val,
+            'progress_bar': tqdm_dict,
+            'log': tqdm_dict
         })
 
         # can also return just a scalar instead of a dict (return loss_val)
         return output
 
-    def validation_step(self, data_batch, batch_i):
+    def validation_step(self, batch, batch_idx):
         """
         Lightning calls this inside the validation loop
-        :param data_batch:
+        :param batch:
         :return:
         """
-        x, y = data_batch
+        x, y = batch
         x = x.view(x.size(0), -1)
         y_hat = self.forward(x)
 
@@ -126,7 +131,7 @@ class LightningTemplateModel(LightningModule):
             val_acc = val_acc.cuda(loss_val.device.index)
 
         # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
+        if self.trainer.use_dp or self.trainer.use_ddp2:
             loss_val = loss_val.unsqueeze(0)
             val_acc = val_acc.unsqueeze(0)
 
@@ -154,21 +159,22 @@ class LightningTemplateModel(LightningModule):
             val_loss = output['val_loss']
 
             # reduce manually when using dp
-            if self.trainer.use_dp:
+            if self.trainer.use_dp or self.trainer.use_ddp2:
                 val_loss = torch.mean(val_loss)
             val_loss_mean += val_loss
 
             # reduce manually when using dp
             val_acc = output['val_acc']
-            if self.trainer.use_dp:
-                val_acc_mean = torch.mean(val_acc)
+            if self.trainer.use_dp or self.trainer.use_ddp2:
+                val_acc = torch.mean(val_acc)
 
-            val_acc_mean += val_acc_mean
+            val_acc_mean += val_acc
 
         val_loss_mean /= len(outputs)
         val_acc_mean /= len(outputs)
-        tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean}
-        return tqdm_dic
+        tqdm_dict = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean}
+        result = {'progress_bar': tqdm_dict, 'log': tqdm_dict, 'val_loss': val_loss_mean}
+        return result
 
     # ---------------------
     # TRAINING SETUP
@@ -189,40 +195,37 @@ class LightningTemplateModel(LightningModule):
         dataset = MNIST(root=self.hparams.data_root, train=train,
                         transform=transform, download=True)
 
-        # when using multi-node we need to add the datasampler
+        # when using multi-node (ddp) we need to add the  datasampler
         train_sampler = None
         batch_size = self.hparams.batch_size
 
-        try:
-            if self.on_gpu:
-                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
-                batch_size = batch_size // self.trainer.world_size  # scale batch size
-        except Exception:
-            pass
+        if self.use_ddp:
+            train_sampler = DistributedSampler(dataset)
 
         should_shuffle = train_sampler is None
         loader = DataLoader(
             dataset=dataset,
             batch_size=batch_size,
             shuffle=should_shuffle,
-            sampler=train_sampler
+            sampler=train_sampler,
+            num_workers=0
         )
 
         return loader
 
     @pl.data_loader
-    def tng_dataloader(self):
-        print('tng data loader called')
+    def train_dataloader(self):
+        logging.info('training data loader called')
         return self.__dataloader(train=True)
 
     @pl.data_loader
     def val_dataloader(self):
-        print('val data loader called')
+        logging.info('val data loader called')
         return self.__dataloader(train=False)
 
     @pl.data_loader
     def test_dataloader(self):
-        print('test data loader called')
+        logging.info('test data loader called')
         return self.__dataloader(train=False)
 
     @staticmethod
@@ -233,31 +236,23 @@ class LightningTemplateModel(LightningModule):
         :param root_dir:
         :return:
         """
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+        parser = ArgumentParser(parents=[parent_parser])
 
         # param overwrites
-        # parser.set_defaults(gradient_clip=5.0)
+        # parser.set_defaults(gradient_clip_val=5.0)
 
         # network params
         parser.add_argument('--in_features', default=28 * 28, type=int)
         parser.add_argument('--out_features', default=10, type=int)
         # use 500 for CPU, 50000 for GPU to see speed difference
         parser.add_argument('--hidden_dim', default=50000, type=int)
-        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
+        parser.add_argument('--drop_prob', default=0.2, type=float)
+        parser.add_argument('--learning_rate', default=0.001, type=float)
 
         # 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 * 8, type=float,
-                        options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
-        parser.opt_list('--optimizer_name', default='adam', type=str,
-                        options=['adam'], tunable=False)
-
-        # if using 2 nodes with 4 gpus each the batch size here
-        #  (256) will be 256 / (2*8) = 16 per gpu
-        parser.opt_list('--batch_size', default=256 * 8, type=int,
-                        options=[32, 64, 128, 256], tunable=False,
-                        help='batch size will be divided over all gpus being used across all nodes')
+        parser.add_argument('--optimizer_name', default='adam', type=str)
+        parser.add_argument('--batch_size', default=64, type=int)
         return parser

pl_examples/domain_templates/gan.py

@@ -0,0 +1,208 @@
+"""
+To run this template just do:
+python gan.py
+
+After a few epochs, launch tensorboard to see the images being generated at every batch.
+
+tensorboard --logdir default
+"""
+import os
+from argparse import ArgumentParser
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+
+import pytorch_lightning as pl
+
+
+class Generator(nn.Module):
+    def __init__(self, latent_dim, img_shape):
+        super(Generator, self).__init__()
+        self.img_shape = img_shape
+
+        def block(in_feat, out_feat, normalize=True):
+            layers = [nn.Linear(in_feat, out_feat)]
+            if normalize:
+                layers.append(nn.BatchNorm1d(out_feat, 0.8))
+            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            return layers
+
+        self.model = nn.Sequential(
+            *block(latent_dim, 128, normalize=False),
+            *block(128, 256),
+            *block(256, 512),
+            *block(512, 1024),
+            nn.Linear(1024, int(np.prod(img_shape))),
+            nn.Tanh()
+        )
+
+    def forward(self, z):
+        img = self.model(z)
+        img = img.view(img.size(0), *self.img_shape)
+        return img
+
+
+class Discriminator(nn.Module):
+    def __init__(self, img_shape):
+        super(Discriminator, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Linear(int(np.prod(img_shape)), 512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(512, 256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(256, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, img):
+        img_flat = img.view(img.size(0), -1)
+        validity = self.model(img_flat)
+
+        return validity
+
+
+class GAN(pl.LightningModule):
+
+    def __init__(self, hparams):
+        super(GAN, self).__init__()
+        self.hparams = hparams
+
+        # networks
+        mnist_shape = (1, 28, 28)
+        self.generator = Generator(latent_dim=hparams.latent_dim, img_shape=mnist_shape)
+        self.discriminator = Discriminator(img_shape=mnist_shape)
+
+        # cache for generated images
+        self.generated_imgs = None
+        self.last_imgs = None
+
+    def forward(self, z):
+        return self.generator(z)
+
+    def adversarial_loss(self, y_hat, y):
+        return F.binary_cross_entropy(y_hat, y)
+
+    def training_step(self, batch, batch_nb, optimizer_i):
+        imgs, _ = batch
+        self.last_imgs = imgs
+
+        # train generator
+        if optimizer_i == 0:
+            # sample noise
+            z = torch.randn(imgs.shape[0], self.hparams.latent_dim)
+
+            # match gpu device (or keep as cpu)
+            if self.on_gpu:
+                z = z.cuda(imgs.device.index)
+
+            # generate images
+            self.generated_imgs = self.forward(z)
+
+            # log sampled images
+            # sample_imgs = self.generated_imgs[:6]
+            # grid = torchvision.utils.make_grid(sample_imgs)
+            # self.logger.experiment.add_image('generated_images', grid, 0)
+
+            # ground truth result (ie: all fake)
+            valid = torch.ones(imgs.size(0), 1)
+
+            # adversarial loss is binary cross-entropy
+            g_loss = self.adversarial_loss(self.discriminator(self.generated_imgs), valid)
+            tqdm_dict = {'g_loss': g_loss}
+            output = OrderedDict({
+                'loss': g_loss,
+                'progress_bar': tqdm_dict,
+                'log': tqdm_dict
+            })
+            return output
+
+        # train discriminator
+        if optimizer_i == 1:
+            # Measure discriminator's ability to classify real from generated samples
+
+            # how well can it label as real?
+            valid = torch.ones(imgs.size(0), 1)
+            real_loss = self.adversarial_loss(self.discriminator(imgs), valid)
+
+            # how well can it label as fake?
+            fake = torch.zeros(imgs.size(0), 1)
+            fake_loss = self.adversarial_loss(
+                self.discriminator(self.generated_imgs.detach()), fake)
+
+            # discriminator loss is the average of these
+            d_loss = (real_loss + fake_loss) / 2
+            tqdm_dict = {'d_loss': d_loss}
+            output = OrderedDict({
+                'loss': d_loss,
+                'progress_bar': tqdm_dict,
+                'log': tqdm_dict
+            })
+            return output
+
+    def configure_optimizers(self):
+        lr = self.hparams.lr
+        b1 = self.hparams.b1
+        b2 = self.hparams.b2
+
+        opt_g = torch.optim.Adam(self.generator.parameters(), lr=lr, betas=(b1, b2))
+        opt_d = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(b1, b2))
+        return [opt_g, opt_d], []
+
+    @pl.data_loader
+    def train_dataloader(self):
+        transform = transforms.Compose([transforms.ToTensor(),
+                                        transforms.Normalize([0.5], [0.5])])
+        dataset = MNIST(os.getcwd(), train=True, download=True, transform=transform)
+        return DataLoader(dataset, batch_size=self.hparams.batch_size)
+
+    def on_epoch_end(self):
+        z = torch.randn(8, self.hparams.latent_dim)
+        # match gpu device (or keep as cpu)
+        if self.on_gpu:
+            z = z.cuda(self.last_imgs.device.index)
+
+        # log sampled images
+        sample_imgs = self.forward(z)
+        grid = torchvision.utils.make_grid(sample_imgs)
+        self.logger.experiment.add_image(f'generated_images', grid, self.current_epoch)
+
+
+def main(hparams):
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = GAN(hparams)
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = pl.Trainer()
+
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument("--batch_size", type=int, default=64, help="size of the batches")
+    parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
+    parser.add_argument("--b1", type=float, default=0.5,
+                        help="adam: decay of first order momentum of gradient")
+    parser.add_argument("--b2", type=float, default=0.999,
+                        help="adam: decay of first order momentum of gradient")
+    parser.add_argument("--latent_dim", type=int, default=100,
+                        help="dimensionality of the latent space")
+
+    hparams = parser.parse_args()
+
+    main(hparams)

pl_examples/multi_node_examples/README.md

@@ -0,0 +1,21 @@
+# Multi-node example   
+
+This demo launches a job using 2 GPUs on 2 different nodes (4 GPUs total).
+To run this demo do the following:
+
+1. Log into the jumphost node of your SLURM-managed cluster.  
+2. Create a conda environment with Lightning and a GPU PyTorch version.   
+3. Choose a script to submit    
+
+#### DDP  
+Submit this job to run with distributedDataParallel (2 nodes, 2 gpus each)
+```bash
+sbatch ddp_job_submit.sh YourEnv
+```
+
+#### DDP2  
+Submit this job to run with a different implementation of distributedDataParallel.
+In this version, each node acts like DataParallel but syncs across nodes like DDP.
+```bash
+sbatch ddp2_job_submit.sh YourEnv
+```

pl_examples/multi_node_examples/ddp2_job_submit.sh

@@ -0,0 +1,27 @@
+#!/bin/bash -l
+
+# SLURM SUBMIT SCRIPT
+#SBATCH --nodes=2
+#SBATCH --gres=gpu:2
+#SBATCH --ntasks-per-node=1
+#SBATCH --mem=0
+#SBATCH --time=0-02:00:00
+
+# activate conda env
+source activate $1
+
+# -------------------------
+# debugging flags (optional)
+ export NCCL_DEBUG=INFO
+ export PYTHONFAULTHANDLER=1
+
+# on your cluster you might need these:
+# set the network interface
+# export NCCL_SOCKET_IFNAME=^docker0,lo
+
+# might need the latest cuda
+# module load NCCL/2.4.7-1-cuda.10.0
+# -------------------------
+
+# run script from above
+srun python3 multi_node_ddp2_demo.py

pl_examples/multi_node_examples/ddp_job_submit.sh

@@ -0,0 +1,27 @@
+#!/bin/bash -l
+
+# SLURM SUBMIT SCRIPT
+#SBATCH --nodes=2
+#SBATCH --gres=gpu:2
+#SBATCH --ntasks-per-node=2
+#SBATCH --mem=0
+#SBATCH --time=0-02:00:00
+
+# activate conda env
+source activate $1
+
+# -------------------------
+# debugging flags (optional)
+ export NCCL_DEBUG=INFO
+ export PYTHONFAULTHANDLER=1
+
+# on your cluster you might need these:
+# set the network interface
+# export NCCL_SOCKET_IFNAME=^docker0,lo
+
+# might need the latest cuda
+# module load NCCL/2.4.7-1-cuda.10.0
+# -------------------------
+
+# run script from above
+srun python3 multi_node_ddp_demo.py

pl_examples/multi_node_examples/multi_node_ddp2_demo.py

@@ -0,0 +1,55 @@
+"""
+Multi-node example (GPU)
+"""
+import os
+from argparse import ArgumentParser
+
+import numpy as np
+import torch
+
+from pl_examples.basic_examples.lightning_module_template import LightningTemplateModel
+from pytorch_lightning import Trainer
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = LightningTemplateModel(hparams)
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        gpus=2,
+        nb_gpu_nodes=2,
+        distributed_backend='ddp2'
+    )
+
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    parent_parser = ArgumentParser(add_help=False)
+
+    # each LightningModule defines arguments relevant to it
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    main(hyperparams)

pl_examples/multi_node_examples/multi_node_ddp_demo.py

@@ -0,0 +1,55 @@
+"""
+Multi-node example (GPU)
+"""
+import os
+from argparse import ArgumentParser
+
+import numpy as np
+import torch
+
+from pl_examples.basic_examples.lightning_module_template import LightningTemplateModel
+from pytorch_lightning import Trainer
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = LightningTemplateModel(hparams)
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        gpus=2,
+        nb_gpu_nodes=2,
+        distributed_backend='ddp'
+    )
+
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    parent_parser = ArgumentParser(add_help=False)
+
+    # each LightningModule defines arguments relevant to it
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    main(hyperparams)

pytorch_lightning/__init__.py

@@ -1,9 +1,36 @@
-from .models.trainer import Trainer
-from .root_module.root_module import LightningModule
-from .root_module.decorators import data_loader
+"""Package info"""
 
-__all__ = [
-    'Trainer',
-    'LightningModule',
-    'data_loader',
-]
+__version__ = '0.5.3'
+__author__ = ' William Falcon et al.'
+__author_email__ = 'waf2107@columbia.edu'
+__license__ = 'Apache-2.0'
+__homepage__ = 'https://github.com/williamFalcon/pytorch-lightning'
+__docs__ = """# PyTorch Lightning
+
+The lightweight PyTorch wrapper for ML researchers. Scale your models. Write less boilerplate.
+"""
+
+
+try:
+    # This variable is injected in the __builtins__ by the build
+    # process. It used to enable importing subpackages of skimage when
+    # the binaries are not built
+    __LIGHTNING_SETUP__
+except NameError:
+    __LIGHTNING_SETUP__ = False
+
+if __LIGHTNING_SETUP__:
+    import sys
+    sys.stderr.write('Partial import of skimage during the build process.\n')
+    # We are not importing the rest of the scikit during the build
+    # process, as it may not be compiled yet
+else:
+    from .trainer.trainer import Trainer
+    from .root_module.root_module import LightningModule
+    from .root_module.decorators import data_loader
+
+    __all__ = [
+        'Trainer',
+        'LightningModule',
+        'data_loader',
+    ]

pytorch_lightning/callbacks/__init__.py

@@ -1,6 +1,7 @@
-from .pt_callbacks import EarlyStopping, ModelCheckpoint
+from .pt_callbacks import EarlyStopping, ModelCheckpoint, GradientAccumulationScheduler
 
 __all__ = [
     'EarlyStopping',
     'ModelCheckpoint',
+    'GradientAccumulationScheduler',
 ]

pytorch_lightning/callbacks/pt_callbacks.py

@@ -1,6 +1,7 @@
 import os
 import shutil
-
+import logging
+import warnings
 import numpy as np
 
 from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
@@ -11,7 +12,6 @@ class Callback(object):
     # Properties
         params: dict. Training parameters
             (eg. verbosity, batch size, number of epochs...).
-        model: instance of `keras.models.Model`.
             Reference of the model being trained.
     The `logs` dictionary that callback methods
     take as argument will contain keys for quantities relevant to
@@ -92,7 +92,7 @@ class EarlyStopping(Callback):
         self.stopped_epoch = 0
 
         if mode not in ['auto', 'min', 'max']:
-            print('EarlyStopping mode %s is unknown, fallback to auto mode.' % mode)
+            logging.info(f'EarlyStopping mode {mode} is unknown, fallback to auto mode.')
             mode = 'auto'
 
         if mode == 'min':
@@ -122,10 +122,12 @@ class EarlyStopping(Callback):
         current = logs.get(self.monitor)
         stop_training = False
         if current is None:
-            print('Early stopping conditioned on metric `%s` '
-                  'which is not available. Available metrics are: %s' %
-                  (self.monitor, ','.join(list(logs.keys()))), RuntimeWarning)
-            exit(-1)
+            warnings.warn(
+                f'Early stopping conditioned on metric `{self.monitor}`'
+                f' which is not available. Available metrics are: {",".join(list(logs.keys()))}',
+                RuntimeWarning)
+            stop_training = True
+            return stop_training
 
         if self.monitor_op(current - self.min_delta, self.best):
             self.best = current
@@ -141,7 +143,7 @@ class EarlyStopping(Callback):
 
     def on_train_end(self, logs=None):
         if self.stopped_epoch > 0 and self.verbose > 0:
-            print('Epoch %05d: early stopping' % (self.stopped_epoch + 1))
+            logging.info(f'Epoch {self.stopped_epoch + 1:05d}: early stopping')
 
 
 class ModelCheckpoint(Callback):
@@ -174,9 +176,19 @@ class ModelCheckpoint(Callback):
     """
 
     def __init__(self, filepath, monitor='val_loss', verbose=0,
-                 save_best_only=False, save_weights_only=False,
+                 save_best_only=True, save_weights_only=False,
                  mode='auto', period=1, prefix=''):
         super(ModelCheckpoint, self).__init__()
+        if (
+            save_best_only and
+            os.path.isdir(filepath) and
+            len(os.listdir(filepath)) > 0
+        ):
+            warnings.warn(
+                f"Checkpoint directory {filepath} exists and is not empty with save_best_only=True."
+                "All files in this directory will be deleted when a checkpoint is saved!"
+            )
+
         self.monitor = monitor
         self.verbose = verbose
         self.filepath = filepath
@@ -187,8 +199,9 @@ class ModelCheckpoint(Callback):
         self.prefix = prefix
 
         if mode not in ['auto', 'min', 'max']:
-            print('ModelCheckpoint mode %s is unknown, '
-                  'fallback to auto mode.' % (mode), RuntimeWarning)
+            warnings.warn(
+                f'ModelCheckpoint mode {mode} is unknown, '
+                'fallback to auto mode.', RuntimeWarning)
             mode = 'auto'
 
         if mode == 'min':
@@ -232,33 +245,66 @@ class ModelCheckpoint(Callback):
             if self.save_best_only:
                 current = logs.get(self.monitor)
                 if current is None:
-                    print('Can save best model only with %s available,'
-                          ' skipping.' % (self.monitor), RuntimeWarning)
+                    warnings.warn(
+                        f'Can save best model only with {self.monitor} available,'
+                        ' skipping.', RuntimeWarning)
                 else:
                     if self.monitor_op(current, self.best):
                         if self.verbose > 0:
-                            print('\nEpoch %05d: %s improved from %0.5f to %0.5f,'
-                                  ' saving model to %s'
-                                  % (epoch + 1, self.monitor, self.best,
-                                     current, filepath))
+                            logging.info(
+                                f'\nEpoch {epoch + 1:05d}: {self.monitor} improved'
+                                f' from {self.best:0.5f} to {current:0.5f},',
+                                f' saving model to {filepath}')
                         self.best = current
                         self.save_model(filepath, overwrite=True)
 
                     else:
                         if self.verbose > 0:
-                            print('\nEpoch %05d: %s did not improve' %
-                                  (epoch + 1, self.monitor))
+                            logging.info(
+                                f'\nEpoch {epoch + 1:05d}: {self.monitor} did not improve')
             else:
                 if self.verbose > 0:
-                    print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
+                    logging.info(f'\nEpoch {epoch + 1:05d}: saving model to {filepath}')
                 self.save_model(filepath, overwrite=False)
 
 
+class GradientAccumulationScheduler(Callback):
+    """Change gradient accumulation factor according to scheduling.
+    # Arguments
+        scheduling: dict, scheduling in format {epoch: accumulation_factor}
+    """
+
+    def __init__(self, scheduling: dict):
+        if scheduling == {}:  # empty dict error
+            raise TypeError("Empty dict cannot be interpreted correct")
+
+        for key in scheduling.keys():
+            if not isinstance(key, int) or not isinstance(scheduling[key], int):
+                raise TypeError("All epoches and accumulation factor must be integers")
+
+        minimal_epoch = min(scheduling.keys())
+        if minimal_epoch < 1:
+            msg = f"Epochs indexing from 1, epoch {minimal_epoch} cannot be interpreted correct"
+            raise IndexError(msg)
+        elif minimal_epoch != 1:  # if user didnt define first epoch accumulation factor
+            scheduling.update({1: 1})
+
+        self.scheduling = scheduling
+        self.epochs = sorted(scheduling.keys())
+
+    def on_epoch_begin(self, epoch, trainer):
+        epoch += 1  # indexing epochs from 1
+        for i in reversed(range(len(self.epochs))):
+            if epoch >= self.epochs[i]:
+                trainer.accumulate_grad_batches = self.scheduling.get(self.epochs[i])
+                break
+
+
 if __name__ == '__main__':
     c = EarlyStopping(min_delta=0.9, patience=2, verbose=True)
     losses = [10, 9, 8, 8, 6, 4.3, 5, 4.4, 2.8, 2.5]
     for i, loss in enumerate(losses):
         should_stop = c.on_epoch_end(i, logs={'val_loss': loss})
-        print(loss)
+        logging.info(loss)
         if should_stop:
             break

pytorch_lightning/logging/__init__.py

@@ -0,0 +1,18 @@
+from os import environ
+from .base import LightningLoggerBase, rank_zero_only
+
+try:
+    from .test_tube_logger import TestTubeLogger
+except ImportError:
+    pass
+try:
+    from .mlflow_logger import MLFlowLogger
+except ImportError:
+    pass
+try:
+    # needed to prevent ImportError and duplicated logs.
+    environ["COMET_DISABLE_AUTO_LOGGING"] = "1"
+
+    from .comet_logger import CometLogger
+except ImportError:
+    del environ["COMET_DISABLE_AUTO_LOGGING"]

pytorch_lightning/logging/base.py

@@ -0,0 +1,76 @@
+from functools import wraps
+
+
+def rank_zero_only(fn):
+    """Decorate a logger method to run it only on the process with rank 0
+
+    :param fn: Function to decorate
+    """
+
+    @wraps(fn)
+    def wrapped_fn(self, *args, **kwargs):
+        if self.rank == 0:
+            fn(self, *args, **kwargs)
+
+    return wrapped_fn
+
+
+class LightningLoggerBase(object):
+    """Base class for experiment loggers"""
+
+    def __init__(self):
+        self._rank = 0
+
+    def log_metrics(self, metrics, step_num):
+        """Record metrics
+
+        :param metric: Dictionary with metric names as keys and measured
+        quanties as values
+        :param step_num: Step number at which the metrics should be recorded
+        """
+        raise NotImplementedError()
+
+    def log_hyperparams(self, params):
+        """Record hyperparameters
+
+        :param params: argparse.Namespace containing the hyperparameters
+        """
+        raise NotImplementedError()
+
+    def save(self):
+        """Save log data"""
+        pass
+
+    def finalize(self, status):
+        """Do any processing that is necessary to finalize an experiment
+
+        :param status: Status that the experiment finished with (e.g. success, failed, aborted)
+        """
+        pass
+
+    def close(self):
+        """Do any cleanup that is necessary to close an experiment"""
+        pass
+
+    @property
+    def rank(self):
+        """
+        Process rank. In general, metrics should only be logged by the process
+        with rank 0
+        """
+        return self._rank
+
+    @rank.setter
+    def rank(self, value):
+        """Set the process rank"""
+        self._rank = value
+
+    @property
+    def name(self):
+        """Return the experiment name"""
+        raise NotImplementedError("Sub-classes must provide a name property")
+
+    @property
+    def version(self):
+        """Return the experiment version"""
+        raise NotImplementedError("Sub-classes must provide a version property")

pytorch_lightning/logging/comet_logger.py

@@ -0,0 +1,25 @@
+try:
+    from comet_ml import Experiment as CometExperiment
+except ImportError:
+    raise ImportError('Missing comet_ml package.')
+
+from .base import LightningLoggerBase, rank_zero_only
+
+
+class CometLogger(LightningLoggerBase):
+    def __init__(self, *args, **kwargs):
+        super(CometLogger, self).__init__()
+        self.experiment = CometExperiment(*args, **kwargs)
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        self.experiment.log_parameters(vars(params))
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step_num):
+        # self.experiment.set_epoch(self, metrics.get('epoch', 0))
+        self.experiment.log_metrics(metrics)
+
+    @rank_zero_only
+    def finalize(self, status):
+        self.experiment.end()

pytorch_lightning/logging/mlflow_logger.py

@@ -0,0 +1,70 @@
+from logging import getLogger
+from time import time
+
+try:
+    import mlflow
+except ImportError:
+    raise ImportError('Missing mlflow package.')
+
+from .base import LightningLoggerBase, rank_zero_only
+
+logger = getLogger(__name__)
+
+
+class MLFlowLogger(LightningLoggerBase):
+    def __init__(self, experiment_name, tracking_uri=None, tags=None):
+        super().__init__()
+        self.experiment = mlflow.tracking.MlflowClient(tracking_uri)
+        self.experiment_name = experiment_name
+        self._run_id = None
+        self.tags = tags
+
+    @property
+    def run_id(self):
+        if self._run_id is not None:
+            return self._run_id
+
+        experiment = self.experiment.get_experiment_by_name(self.experiment_name)
+        if experiment is None:
+            logger.warning(
+                f"Experiment with name f{self.experiment_name} not found. Creating it."
+            )
+            self.experiment.create_experiment(self.experiment_name)
+            experiment = self.experiment.get_experiment_by_name(self.experiment_name)
+
+        run = self.experiment.create_run(experiment.experiment_id, tags=self.tags)
+        self._run_id = run.info.run_id
+        return self._run_id
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        for k, v in vars(params).items():
+            self.experiment.log_param(self.run_id, k, v)
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step_num=None):
+        timestamp_ms = int(time() * 1000)
+        for k, v in metrics.items():
+            if isinstance(v, str):
+                logger.warning(
+                    f"Discarding metric with string value {k}={v}"
+                )
+                continue
+            self.experiment.log_metric(self.run_id, k, v, timestamp_ms, step_num)
+
+    def save(self):
+        pass
+
+    @rank_zero_only
+    def finalize(self, status="FINISHED"):
+        if status == 'success':
+            status = 'FINISHED'
+        self.experiment.set_terminated(self.run_id, status)
+
+    @property
+    def name(self):
+        return self.experiment_name
+
+    @property
+    def version(self):
+        return self._run_id

pytorch_lightning/logging/test_tube_logger.py

@@ -0,0 +1,109 @@
+try:
+    from test_tube import Experiment
+except ImportError:
+    raise ImportError('Missing test-tube package.')
+
+from .base import LightningLoggerBase, rank_zero_only
+
+
+class TestTubeLogger(LightningLoggerBase):
+    __test__ = False
+
+    def __init__(
+            self, save_dir, name="default", description=None, debug=False,
+            version=None, create_git_tag=False
+    ):
+        super().__init__()
+        self.save_dir = save_dir
+        self._name = name
+        self.description = description
+        self.debug = debug
+        self._version = version
+        self.create_git_tag = create_git_tag
+        self._experiment = None
+
+    @property
+    def experiment(self):
+        if self._experiment is not None:
+            return self._experiment
+
+        self._experiment = Experiment(
+            save_dir=self.save_dir,
+            name=self._name,
+            debug=self.debug,
+            version=self.version,
+            description=self.description,
+            create_git_tag=self.create_git_tag,
+            rank=self.rank,
+        )
+        return self._experiment
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
+        self.experiment.argparse(params)
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step_num=None):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
+        self.experiment.log(metrics, global_step=step_num)
+
+    @rank_zero_only
+    def save(self):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
+        self.experiment.save()
+
+    @rank_zero_only
+    def finalize(self, status):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
+        self.save()
+        self.close()
+
+    @rank_zero_only
+    def close(self):
+        # TODO: HACK figure out where this is being set to true
+        self.experiment.debug = self.debug
+        exp = self.experiment
+        exp.close()
+
+    @property
+    def rank(self):
+        return self._rank
+
+    @rank.setter
+    def rank(self, value):
+        self._rank = value
+        if self._experiment is not None:
+            self.experiment.rank = value
+
+    @property
+    def name(self):
+        if self._experiment is None:
+            return self._name
+        else:
+            return self.experiment.name
+
+    @property
+    def version(self):
+        if self._experiment is None:
+            return self._version
+        else:
+            return self.experiment.version
+
+    # Test tube experiments are not pickleable, so we need to override a few
+    # methods to get DDP working. See
+    # https://docs.python.org/3/library/pickle.html#handling-stateful-objects
+    # for more info.
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["_experiment"] = self.experiment.get_meta_copy()
+        return state
+
+    def __setstate__(self, state):
+        self._experiment = state["_experiment"].get_non_ddp_exp()
+        del state["_experiment"]
+        self.__dict__.update(state)

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -1,11 +1,11 @@
-from torch.nn import DataParallel
-from torch.nn.parallel import DistributedDataParallel
 import itertools
+import threading
 from itertools import chain
 
-import threading
 import torch
 from torch.cuda._utils import _get_device_index
+from torch.nn import DataParallel
+from torch.nn.parallel import DistributedDataParallel
 
 
 def _find_tensors(obj):  # pragma: no cover
@@ -56,6 +56,8 @@ class LightningDataParallel(DataParallel):
             # lightning
             if self.module.training:
                 return self.module.training_step(*inputs[0], **kwargs[0])
+            elif self.module.testing:
+                return self.module.test_step(*inputs[0], **kwargs[0])
             else:
                 return self.module.validation_step(*inputs[0], **kwargs[0])
 
@@ -85,16 +87,18 @@ class LightningDistributedDataParallel(DistributedDataParallel):
                 # --------------
                 # normal
                 # output = self.module(*inputs[0], **kwargs[0])
-
                 # lightning
                 if self.module.training:
                     output = self.module.training_step(*inputs[0], **kwargs[0])
+                elif self.module.testing:
+                    output = self.module.test_step(*inputs[0], **kwargs[0])
                 else:
                     output = self.module.validation_step(*inputs[0], **kwargs[0])
             else:
                 outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
                 output = self.gather(outputs, self.output_device)
         else:
+            # normal
             output = self.module(*inputs, **kwargs)
 
         if torch.is_grad_enabled():
@@ -153,6 +157,10 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):  # pragma: n
                 # CHANGE
                 if module.training:
                     output = module.training_step(*input, **kwargs)
+
+                elif module.testing:
+                    output = module.test_step(*input, **kwargs)
+
                 else:
                     output = module.validation_step(*input, **kwargs)
                 # ---------------
@@ -163,6 +171,14 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):  # pragma: n
             with lock:
                 results[i] = e
 
+    # TODO: fix hack (maybe not a hack)
+    # make sure each module knows what training state it's in...
+    # fixes weird bug where copies are out of sync
+    root_m = modules[0]
+    for m in modules[1:]:
+        m.training = root_m.training
+        m.testing = root_m.testing
+
     if len(modules) > 1:
         threads = [threading.Thread(target=_worker,
                                     args=(i, module, input, kwargs, device))

pytorch_lightning/root_module/decorators.py

@@ -1,3 +1,5 @@
+import traceback
+
 
 def data_loader(fn):
     """
@@ -8,10 +10,24 @@ def data_loader(fn):
 
     attr_name = '_lazy_' + fn.__name__
 
-    @property
-    def _data_loader(self):
-        if not hasattr(self, attr_name):
-            setattr(self, attr_name, fn(self))
-        return getattr(self, attr_name)
+    def _get_data_loader(self):
+        try:
+            value = getattr(self, attr_name)
+        except AttributeError:
+            try:
+                value = fn(self)  # Lazy evaluation, done only once.
+                if (
+                        value is not None and
+                        not isinstance(value, list) and
+                        fn.__name__ in ['test_dataloader', 'val_dataloader']
+                ):
+                    value = [value]
+            except AttributeError as e:
+                # Guard against AttributeError suppression. (Issue #142)
+                traceback.print_exc()
+                error = f'{fn.__name__}: An AttributeError was encountered: ' + str(e)
+                raise RuntimeError(error) from e
+            setattr(self, attr_name, value)  # Memoize evaluation.
+        return value
 
-    return _data_loader
+    return _get_data_loader

pytorch_lightning/root_module/hooks.py

@@ -1,16 +1,24 @@
 import torch
 
 
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+
+
 class ModelHooks(torch.nn.Module):
 
     def on_sanity_check_start(self):
         """
-        Called before starting validate
+        Called before starting evaluate
         :return:
         """
         pass
 
-    def on_batch_start(self, data_batch):
+    def on_batch_start(self, batch):
         pass
 
     def on_batch_end(self):
@@ -28,9 +36,6 @@ class ModelHooks(torch.nn.Module):
     def on_post_performance_check(self):
         pass
 
-    def on_tng_metrics(self, metrics):
-        pass
-
     def on_before_zero_grad(self, optimizer):
         """
         Called after optimizer.step() and before optimizer.zero_grad()
@@ -51,3 +56,17 @@ class ModelHooks(torch.nn.Module):
         :return:
         """
         pass
+
+    def backward(self, use_amp, loss, optimizer):
+        """
+        Override backward with your own implementation if you need to
+        :param use_amp: Whether amp was requested or not
+        :param loss: Loss is already scaled by accumulated grads
+        :param optimizer: Current optimizer being used
+        :return:
+        """
+        if use_amp:
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+        else:
+            loss.backward()

pytorch_lightning/root_module/memory.py

@@ -3,20 +3,23 @@ Generates a summary of a model's layers and dimensionality
 '''
 
 import gc
-
-import torch
+import os
 import subprocess
+
 import numpy as np
 import pandas as pd
+import torch
+import logging
 
 
 class ModelSummary(object):
 
-    def __init__(self, model):
+    def __init__(self, model, mode='full'):
         '''
         Generates summaries of model layers and dimensions.
         '''
         self.model = model
+        self.mode = mode
         self.in_sizes = []
         self.out_sizes = []
 
@@ -28,9 +31,20 @@ class ModelSummary(object):
     def __repr__(self):
         return self.summary.__str__()
 
+    def named_modules(self):
+        if self.mode == 'full':
+            mods = self.model.named_modules()
+            mods = list(mods)[1:]  # do not include root module (LightningModule)
+        elif self.mode == 'top':
+            # the children are the top-level modules
+            mods = self.model.named_children()
+        else:
+            mods = []
+        return list(mods)
+
     def get_variable_sizes(self):
         '''Run sample input through each layer to get output sizes'''
-        mods = list(self.model.modules())
+        mods = self.named_modules()
         in_sizes = []
         out_sizes = []
         input_ = self.model.example_input_array
@@ -43,8 +57,7 @@ class ModelSummary(object):
 
         with torch.no_grad():
 
-            for i in range(1, len(mods)):
-                m = mods[i]
+            for _, m in mods:
                 if type(input_) is list or type(input_) is tuple:  # pragma: no cover
                     out = m(*input_)
                 else:
@@ -72,16 +85,17 @@ class ModelSummary(object):
 
         self.in_sizes = in_sizes
         self.out_sizes = out_sizes
+        assert len(in_sizes) == len(out_sizes)
         return
 
     def get_layer_names(self):
         '''Collect Layer Names'''
-        mods = list(self.model.named_modules())
+        mods = self.named_modules()
         names = []
         layers = []
-        for m in mods[1:]:
-            names += [m[0]]
-            layers += [str(m[1].__class__)]
+        for name, m in mods:
+            names += [name]
+            layers += [str(m.__class__)]
 
         layer_types = [x.split('.')[-1][:-2] for x in layers]
 
@@ -91,11 +105,9 @@ class ModelSummary(object):
 
     def get_parameter_sizes(self):
         '''Get sizes of all parameters in `model`'''
-        mods = list(self.model.modules())
+        mods = self.named_modules()
         sizes = []
-
-        for i in range(1, len(mods)):
-            m = mods[i]
+        for _, m in mods:
             p = list(m.parameters())
             modsz = []
             for j in range(len(p)):
@@ -133,9 +145,9 @@ class ModelSummary(object):
         df['Name'] = self.layer_names
         df['Type'] = self.layer_types
         df['Params'] = self.param_nums
+        df['Params'] = df['Params'].map(get_human_readable_count)
 
         if self.model.example_input_array is not None:
-
             df['In_sizes'] = self.in_sizes
             df['Out_sizes'] = self.out_sizes
 
@@ -156,7 +168,7 @@ def print_mem_stack():  # pragma: no cover
     for obj in gc.get_objects():
         try:
             if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
-                print(type(obj), obj.size())
+                logging.info(type(obj), obj.size())
         except Exception:
             pass
 
@@ -178,6 +190,24 @@ def count_mem_items():  # pragma: no cover
     return nb_params, nb_tensors
 
 
+def get_memory_profile(mode):
+    """
+    'all' means return memory for all gpus
+    'min_max' means return memory for max and min
+    :param mode:
+    :return:
+    """
+    memory_map = get_gpu_memory_map()
+
+    if mode == 'min_max':
+        min_index, min_memory = min(memory_map.items(), key=lambda item: item[1])
+        max_index, max_memory = max(memory_map.items(), key=lambda item: item[1])
+
+        memory_map = {min_index: min_memory, max_index: max_memory}
+
+    return memory_map
+
+
 def get_gpu_memory_map():
     """Get the current gpu usage.
 
@@ -187,15 +217,41 @@ def get_gpu_memory_map():
         Keys are device ids as integers.
         Values are memory usage as integers in MB.
     """
-    result = subprocess.check_output(
+    result = subprocess.run(
         [
-            'nvidia-smi', '--query-gpu=memory.used',
-            '--format=csv,nounits,noheader'
-        ], encoding='utf-8')
+            'nvidia-smi',
+            '--query-gpu=memory.used',
+            '--format=csv,nounits,noheader',
+        ],
+        encoding='utf-8',
+        capture_output=True,
+        check=True)
     # Convert lines into a dictionary
-    gpu_memory = [int(x) for x in result.strip().split('\n')]
-    gpu_memory_map = {}
-    for k, v in zip(range(len(gpu_memory)), gpu_memory):
-        k = 'gpu_%i' % k
-        gpu_memory_map[k] = v
+    gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)]
+    gpu_memory_map = {f'gpu_{index}': memory for index, memory in enumerate(gpu_memory)}
     return gpu_memory_map
+
+
+def get_human_readable_count(number):
+    """
+    Abbreviates an integer number with K, M, B, T for thousands, millions,
+    billions and trillions, respectively.
+    Examples:
+        123     -> 123
+        1234    -> 1 K       (one thousand)
+        2e6     -> 2 M       (two million)
+        3e9     -> 3 B       (three billion)
+        4e12    -> 4 T       (four trillion)
+        5e15    -> 5,000 T
+    :param number: a positive integer number
+    :returns a string formatted according to the pattern described above.
+    """
+    assert number >= 0
+    labels = [' ', 'K', 'M', 'B', 'T']
+    num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1)
+    num_groups = int(np.ceil(num_digits / 3))
+    num_groups = min(num_groups, len(labels))  # don't abbreviate beyond trillions
+    shift = -3 * (num_groups - 1)
+    number = number * (10 ** shift)
+    index = num_groups - 1
+    return f'{int(number):,d} {labels[index]}'

pytorch_lightning/root_module/model_saving.py

@@ -1,12 +1,3 @@
-import os
-import re
-
-import torch
-
-from pytorch_lightning.pt_overrides.override_data_parallel import (
-    LightningDistributedDataParallel, LightningDataParallel)
-
-
 class ModelIO(object):
 
     def on_load_checkpoint(self, checkpoint):
@@ -41,221 +32,3 @@ class ModelIO(object):
         :return:
         """
         pass
-
-
-class TrainerIO(object):
-
-    def __get_model(self):
-        is_dp_module = isinstance(self.model, (LightningDistributedDataParallel,
-                                               LightningDataParallel))
-        model = self.model.module if is_dp_module else self.model
-        return model
-
-    # --------------------
-    # MODEL SAVE CHECKPOINT
-    # --------------------
-    def save_checkpoint(self, filepath):
-        checkpoint = self.dump_checkpoint()
-
-        # do the actual save
-        torch.save(checkpoint, filepath)
-
-    def restore(self, checkpoint_path, on_gpu):
-
-        if on_gpu:
-            checkpoint = torch.load(checkpoint_path)
-        else:
-            checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
-
-        # load training state (affects trainer only)
-        self.restore_training_state(checkpoint)
-
-        # load model state
-        model = self.__get_model()
-
-        # load the state_dict on the model automatically
-        model.load_state_dict(checkpoint['state_dict'])
-
-    def dump_checkpoint(self):
-
-        checkpoint = {
-            'epoch': self.current_epoch,
-            'global_step': self.global_step
-        }
-
-        if self.checkpoint_callback is not None:
-            checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
-
-        if self.early_stop_callback is not None:
-            checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
-            checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
-
-        # save optimizers
-        optimizer_states = []
-        for i, optimizer in enumerate(self.optimizers):
-            optimizer_states.append(optimizer.state_dict())
-
-        checkpoint['optimizer_states'] = optimizer_states
-
-        # save lr schedulers
-        lr_schedulers = []
-        for i, scheduler in enumerate(self.lr_schedulers):
-            lr_schedulers.append(scheduler.state_dict())
-
-        checkpoint['lr_schedulers'] = lr_schedulers
-
-        # add the state_dict from the model
-        model = self.__get_model()
-        checkpoint['state_dict'] = model.state_dict()
-
-        # give the model a chance to add a few things
-        model.on_save_checkpoint(checkpoint)
-
-        return checkpoint
-
-    # --------------------
-    # HPC IO
-    # --------------------
-    def enable_auto_hpc_walltime_manager(self):
-        if self.cluster is None:
-            return
-
-        # allow test tube to handle model check pointing automatically
-        # only if proc 0 so we don't trigger world_size resubmits
-        if self.proc_rank == 0:
-            self.cluster.set_checkpoint_save_function(
-                self.hpc_save,
-                kwargs={
-                    'folderpath': self.checkpoint_callback.filepath,
-                    'experiment': self.experiment
-                }
-            )
-
-        self.cluster.set_checkpoint_load_function(
-            self.hpc_load,
-            kwargs={
-                'folderpath': self.checkpoint_callback.filepath,
-                'on_gpu': self.on_gpu
-            }
-        )
-
-    def restore_training_state(self, checkpoint):
-        """
-        Restore trainer state.
-        Model will get its change to update
-        :param checkpoint:
-        :return:
-        """
-        if self.checkpoint_callback is not None:
-            self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
-
-        if self.early_stop_callback is not None:
-            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']
-        for optimizer, opt_state in zip(self.optimizers, optimizer_states):
-            optimizer.load_state_dict(opt_state)
-
-        # restore the lr schedulers
-        lr_schedulers = checkpoint['lr_schedulers']
-        for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers):
-            scheduler.load_state_dict(lrs_state)
-
-    # ----------------------------------
-    # 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()
-
-        # close experiment to avoid issues
-        experiment.close()
-
-        ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
-
-        if not os.path.exists(folderpath):
-            os.makedirs(folderpath, exist_ok=True)
-        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
-
-        # give model a chance to do something on hpc_save
-        model = self.__get_model()
-        checkpoint = self.dump_checkpoint()
-
-        model.on_hpc_save(checkpoint)
-
-        # do the actual save
-        torch.save(checkpoint, filepath)
-
-        return filepath
-
-    def hpc_load(self, folderpath, on_gpu):
-        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
-
-        if on_gpu:
-            checkpoint = torch.load(filepath)
-        else:
-            checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
-
-        # load training state (affects trainer only)
-        self.restore_training_state(checkpoint)
-
-        # load model state
-        model = self.__get_model()
-
-        # load the state_dict on the model automatically
-        model.load_state_dict(checkpoint['state_dict'])
-
-        # call model hook
-        model.on_hpc_load(checkpoint)
-
-    def max_ckpt_in_folder(self, path, name_key='ckpt_'):
-        files = os.listdir(path)
-        files = [x for x in files if name_key in x]
-        if len(files) == 0:
-            return 0
-
-        ckpt_vs = []
-        for name in files:
-            name = name.split(name_key)[-1]
-            name = re.sub('[^0-9]', '', name)
-            ckpt_vs.append(int(name))
-
-        return max(ckpt_vs)
-
-
-def load_hparams_from_tags_csv(tags_csv):
-    from argparse import Namespace
-    import pandas as pd
-
-    tags_df = pd.read_csv(tags_csv)
-    dic = tags_df.to_dict(orient='records')
-
-    ns_dict = {row['key']: convert(row['value']) for row in dic}
-
-    ns = Namespace(**ns_dict)
-    return ns
-
-
-def convert(val):
-    constructors = [int, float, str]
-
-    if type(val) is str:
-        if val.lower() == 'true':
-            return True
-        if val.lower() == 'false':
-            return False
-
-    for c in constructors:
-        try:
-            return c(val)
-        except ValueError:
-            pass
-    return val

pytorch_lightning/root_module/root_module.py

@@ -1,10 +1,19 @@
+import os
+import warnings
+import collections
+from argparse import Namespace
+
 import torch
+import torch.distributed as dist
 
-from pytorch_lightning.root_module.memory import ModelSummary
-from pytorch_lightning.root_module.grads import GradInformation
-from pytorch_lightning.root_module.model_saving import ModelIO, load_hparams_from_tags_csv
-from pytorch_lightning.root_module.hooks import ModelHooks
 from pytorch_lightning.root_module.decorators import data_loader
+from pytorch_lightning.root_module.grads import GradInformation
+from pytorch_lightning.root_module.hooks import ModelHooks
+from pytorch_lightning.root_module.memory import ModelSummary
+from pytorch_lightning.root_module.model_saving import ModelIO
+from pytorch_lightning.trainer.trainer_io import load_hparams_from_tags_csv
+import logging
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 
 
 class LightningModule(GradInformation, ModelIO, ModelHooks):
@@ -18,11 +27,15 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         self.global_step = 0
         self.loaded_optimizer_states_dict = {}
         self.trainer = None
-        self.experiment = None
+        self.logger = None
         self.example_input_array = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
+        self.use_dp = False
+        self.use_ddp = False
+        self.use_ddp2 = False
+        self.use_amp = False
 
     def forward(self, *args, **kwargs):
         """
@@ -33,29 +46,127 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         raise NotImplementedError
 
-    def validation_step(self, data_batch, batch_nb):
+    def training_step(self, *args, **kwargs):
         """
-        return whatever outputs will need to be aggregated in validation_end
-        :param data_batch:
-        :return:
+        return loss, dict with metrics for tqdm
+        :param called with batch, batch_nb
+        additional: optimizer_i if multiple optimizers used
+        :return: dict with loss key and optional log, progress keys
+                if implementing training_step, return whatever you need in that step
         """
         raise NotImplementedError
 
+    def training_end(self, *args, **kwargs):
+        """
+        return loss, dict with metrics for tqdm
+        :param called with outputs of training_step
+        :return: dict with loss key and optional log, progress keys
+        """
+        pass
+
+    def validation_step(self, *args, **kwargs):
+        """
+        return whatever outputs will need to be aggregated in validation_end
+        OPTIONAL
+        :param called with batch, batch_nb
+        additional: dataset_i if multiple val datasets used
+        :return:
+        """
+        pass
+
+    def test_step(self, *args, **kwargs):
+        """
+        return whatever outputs will need to be aggregated in test_end
+        OPTIONAL
+        :param called with batch, batch_nb
+        additional: dataset_i if multiple val datasets used
+        :return:
+        """
+        pass
+
     def validation_end(self, outputs):
         """
         Outputs has the appended output after each validation step
+        OPTIONAL
         :param outputs:
         :return: dic_with_metrics for tqdm
         """
-        raise NotImplementedError
+        pass
 
-    def training_step(self, data_batch, batch_nb):
+    def test_end(self, outputs):
         """
-        return loss, dict with metrics for tqdm
-        :param data_batch:
-        :return:
+        Outputs has the appended output after each test step
+        OPTIONAL
+        :param outputs:
+        :return: dic_with_metrics for tqdm
         """
-        raise NotImplementedError
+        pass
+
+    def configure_ddp(self, model, device_ids):
+        """
+        Override to init DDP in a different way or use your own wrapper.
+        Must return model.
+        :param model:
+        :param device_ids:
+        :return: DDP wrapped model
+        """
+        model = LightningDistributedDataParallel(
+            model,
+            device_ids=device_ids,
+            find_unused_parameters=True
+        )
+        return model
+
+    def init_ddp_connection(self, proc_rank, world_size):
+        """
+        Connect all procs in the world using the env:// init
+        Use the first node as the root address
+        """
+
+        # use slurm job id for the port number
+        # guarantees unique ports across jobs from same grid search
+        try:
+            # use the last 4 numbers in the job id as the id
+            default_port = os.environ['SLURM_JOB_ID']
+            default_port = default_port[-4:]
+
+            # all ports should be in the 10k+ range
+            default_port = int(default_port) + 15000
+
+        except Exception as e:
+            default_port = 12910
+
+        # if user gave a port number, use that one instead
+        try:
+            default_port = os.environ['MASTER_PORT']
+        except Exception:
+            os.environ['MASTER_PORT'] = str(default_port)
+
+        # figure out the root node addr
+        try:
+            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        except Exception:
+            root_node = '127.0.0.2'
+
+        root_node = self.trainer.resolve_root_node_address(root_node)
+        os.environ['MASTER_ADDR'] = root_node
+        dist.init_process_group('nccl', rank=proc_rank, world_size=world_size)
+
+    def configure_apex(self, amp, model, optimizers, amp_level):
+        """
+        Override to init AMP your own way
+        Must return a model and list of optimizers
+        :param amp:
+        :param model:
+        :param optimizers:
+        :param amp_level:
+        :return: Apex wrapped model and optimizers
+        """
+        model, optimizers = amp.initialize(
+            model, optimizers, opt_level=amp_level,
+        )
+
+        return model, optimizers
 
     def configure_optimizers(self):
         """
@@ -64,50 +175,108 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         """
         raise NotImplementedError
 
+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, second_order_closure=None):
+        """
+        Do something instead of the standard optimizer behavior
+        :param epoch_nb:
+        :param batch_nb:
+        :param optimizer:
+        :param optimizer_i:
+        :param second_order_closure: closure for second order methods
+        :return:
+        """
+        if isinstance(optimizer, torch.optim.LBFGS):
+            optimizer.step(second_order_closure)
+        else:
+            optimizer.step()
+
+        # clear gradients
+        optimizer.zero_grad()
+
+    def tbptt_split_batch(self, batch, split_size):
+        """
+        Return list of batch splits. Each split will be passed to forward_step to enable truncated
+        back propagation through time. The default implementation splits root level Tensors and
+        Sequences at dim=1 (i.e. time dim). It assumes that each time dim is the same length.
+        :return:
+        """
+        time_dims = [len(x[0]) for x in batch if isinstance(
+            x, torch.Tensor) or isinstance(x, collections.Sequence)]
+        assert len(time_dims) >= 1, "Unable to determine batch time dimension"
+        assert all(x == time_dims[0] for x in time_dims), "Batch time dimension length is ambiguous"
+
+        splits = []
+        for t in range(0, time_dims[0], split_size):
+            batch_split = []
+            for i, x in enumerate(batch):
+                if isinstance(x, torch.Tensor):
+                    split_x = x[:, t:t + split_size]
+                elif isinstance(x, collections.Sequence):
+                    split_x = [None] * len(x)
+                    for batch_idx in range(len(x)):
+                        split_x[batch_idx] = x[batch_idx][t:t + split_size]
+
+                batch_split.append(split_x)
+
+            splits.append(batch_split)
+
+        return splits
+
     @data_loader
     def tng_dataloader(self):
         """
-        Implement a function to load an h5py of this data
+        Implement a PyTorch DataLoader
+        * Deprecated in v0.5.0. use train_dataloader instead. *
         :return:
         """
         raise NotImplementedError
 
+    @data_loader
+    def train_dataloader(self):
+        """
+        Implement a PyTorch DataLoader
+        :return:
+        """
+        #
+        try:
+            output = self.tng_dataloader()
+            warnings.warn("tng_dataloader has been renamed to train_dataloader since v0.5.0",
+                          DeprecationWarning)
+            return output
+        except NotImplementedError:
+            raise NotImplementedError
+
     @data_loader
     def test_dataloader(self):
         """
-        Implement a function to load an h5py of this data
+        Implement a PyTorch DataLoader
         :return:
         """
-        raise NotImplementedError
+        return None
 
     @data_loader
     def val_dataloader(self):
         """
-        Implement a function to load an h5py of this data
+        Implement a PyTorch DataLoader
         :return:
         """
-        raise NotImplementedError
+        return None
 
     @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
+    def load_from_metrics(cls, weights_path, tags_csv):
         """
         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)
+        hparams.__setattr__('on_gpu', False)
 
-        if on_gpu:
-            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)
+        # load on CPU only to avoid OOM issues
+        # then its up to user to put back on GPUs
+        checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 
         # load the state_dict on the model automatically
         model = cls(hparams)
@@ -118,14 +287,48 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
 
         return model
 
-    def summarize(self):
-        model_summary = ModelSummary(self)
-        print(model_summary)
+    @classmethod
+    def load_from_checkpoint(cls, checkpoint_path):
+        """
+        Primary way of loading model from a checkpoint
+        :param checkpoint_path:
+        :param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
+        :return:
+        """
+
+        # load on CPU only to avoid OOM issues
+        # then its up to user to put back on GPUs
+        checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
+        try:
+            ckpt_hparams = checkpoint['hparams']
+        except KeyError:
+            raise IOError(
+                "Checkpoint does not contain hyperparameters. Are your model hyperparameters stored"
+                "in self.hparams?"
+            )
+        hparams = Namespace(**ckpt_hparams)
+
+        # load the state_dict on the model automatically
+        model = cls(hparams)
+        model.load_state_dict(checkpoint['state_dict'])
+
+        # give model a chance to load something
+        model.on_load_checkpoint(checkpoint)
+
+        return model
+
+    def summarize(self, mode):
+        model_summary = ModelSummary(self, mode=mode)
+        logging.info(model_summary)
 
     def freeze(self):
         for param in self.parameters():
             param.requires_grad = False
 
+        self.eval()
+
     def unfreeze(self):
         for param in self.parameters():
             param.requires_grad = True
+
+        self.train()

pytorch_lightning/testing/__init__.py

@@ -0,0 +1,12 @@
+from .lm_test_module import LightningTestModel
+from .lm_test_module_base import LightningTestModelBase
+from .lm_test_module_mixins import (
+    LightningValidationStepMixin,
+    LightningValidationMixin,
+    LightningValidationStepMultipleDataloadersMixin,
+    LightningValidationMultipleDataloadersMixin,
+    LightningTestStepMixin,
+    LightningTestMixin,
+    LightningTestStepMultipleDataloadersMixin,
+    LightningTestMultipleDataloadersMixin,
+)

pytorch_lightning/testing/lm_test_module.py

@@ -1,269 +1,13 @@
-import os
-from collections import OrderedDict
-
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import optim
-from torch.utils.data import DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from torchvision.datasets import MNIST
-from torchvision import transforms
-from test_tube import HyperOptArgumentParser
 
-from pytorch_lightning.root_module.root_module import LightningModule
-from pytorch_lightning import data_loader
+from .lm_test_module_base import LightningTestModelBase
+from .lm_test_module_mixins import LightningValidationMixin, LightningTestMixin
 
 
-class LightningTestModel(LightningModule):
+class LightningTestModel(LightningValidationMixin, LightningTestMixin, LightningTestModelBase):
     """
-    Sample model to show how to define a template
+    Most common test case. Validation and test dataloaders
     """
 
-    def __init__(self, hparams, force_remove_distributed_sampler=False):
-        """
-        Pass in parsed HyperOptArgumentParser to the model
-        :param hparams:
-        """
-        # init superclass
-        super(LightningTestModel, self).__init__()
-        self.hparams = hparams
-
-        self.batch_size = hparams.batch_size
-
-        # if you specify an example input, the summary will show input/output for each layer
-        self.example_input_array = torch.rand(5, 28 * 28)
-
-        # remove to test warning for dist sampler
-        self.force_remove_distributed_sampler = force_remove_distributed_sampler
-
-        # 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):
-        """
-        No special modification required for lightning, define as you normally would
-        :param x:
-        :return:
-        """
-
-        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):
-        """
-        Lightning calls this 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)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if self.trainer.batch_nb % 1 == 0:
-            output = OrderedDict({
-                'loss': loss_val,
-                'prog': {'some_val': loss_val * loss_val}
-            })
-            return output
-        if self.trainer.batch_nb % 2 == 0:
-            return loss_val
-
-    def validation_step(self, data_batch, batch_i):
-        """
-        Lightning calls this 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)
-        val_acc = torch.tensor(val_acc)
-
-        if self.on_gpu:
-            val_acc = val_acc.cuda(loss_val.device.index)
-
-        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
-        if self.trainer.use_dp:
-            loss_val = loss_val.unsqueeze(0)
-            val_acc = val_acc.unsqueeze(0)
-
-        # alternate possible outputs to test
-        if batch_i % 1 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-            })
-            return output
-        if batch_i % 2 == 0:
-            return val_acc
-
-        if batch_i % 3 == 0:
-            output = OrderedDict({
-                'val_loss': loss_val,
-                'val_acc': val_acc,
-                'test_dic': {'val_loss_a': loss_val}
-            })
-            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:
-        """
-        # if returned a scalar from validation_step, outputs is a list of tensor scalars
-        # we return just the average in this case (if we want)
-        # return torch.stack(outputs).mean()
-
-        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 on_tng_metrics(self, logs):
+    def on_training_metrics(self, logs):
         logs['some_tensor_to_test'] = torch.rand(1)
-
-    # ---------------------
-    # TRAINING SETUP
-    # ---------------------
-    def configure_optimizers(self):
-        """
-        return whatever optimizers we want here
-        :return: list of optimizers
-        """
-        # try no scheduler for this model (testing purposes)
-        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
-
-        # test returning only 1 list instead of 2
-        return [optimizer]
-
-    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)
-
-        # when using multi-node we need to add the datasampler
-        train_sampler = None
-        batch_size = self.hparams.batch_size
-
-        try:
-            if self.on_gpu and not self.force_remove_distributed_sampler:
-                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
-                batch_size = batch_size // self.trainer.world_size  # scale batch size
-        except Exception:
-            pass
-
-        should_shuffle = train_sampler is None
-        loader = DataLoader(
-            dataset=dataset,
-            batch_size=batch_size,
-            shuffle=should_shuffle,
-            sampler=train_sampler
-        )
-
-        return loader
-
-    @data_loader
-    def tng_dataloader(self):
-        return self.__dataloader(train=True)
-
-    @data_loader
-    def val_dataloader(self):
-        return self.__dataloader(train=False)
-
-    @data_loader
-    def test_dataloader(self):
-        return self.__dataloader(train=False)
-
-    @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):  # pragma: no cover
-        """
-        Parameters you define here will be available to your model through self.hparams
-        :param parent_parser:
-        :param root_dir:
-        :return:
-        """
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
-
-        # param overwrites
-        # 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, type=int)
-        parser.add_argument('--out_features', default=10, type=int)
-        # use 500 for CPU, 50000 for GPU to see speed difference
-        parser.add_argument('--hidden_dim', default=50000, type=int)
-
-        # 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 * 8, type=float,
-                        options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
-        parser.opt_list('--optimizer_name', default='adam', type=str,
-                        options=['adam'], tunable=False)
-
-        # if using 2 nodes with 4 gpus each the batch size here
-        #  (256) will be 256 / (2*8) = 16 per gpu
-        parser.opt_list('--batch_size', default=256 * 8, type=int,
-                        options=[32, 64, 128, 256], tunable=False,
-                        help='batch size will be divided over all gpus being used across all nodes')
-        return parser

pytorch_lightning/testing/lm_test_module_base.py

@@ -0,0 +1,203 @@
+import os
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torchvision import transforms
+from torchvision.datasets import MNIST
+try:
+    from test_tube import HyperOptArgumentParser
+except ImportError:
+    # TODO: this should be discussed and moved out of this package
+    raise ImportError('Missing test-tube package.')
+
+from pytorch_lightning import data_loader
+from pytorch_lightning.root_module.root_module import LightningModule
+
+
+class LightningTestModelBase(LightningModule):
+    """
+    Base LightningModule for testing. Implements only the required
+    interface
+    """
+
+    def __init__(self, hparams, force_remove_distributed_sampler=False):
+        """
+        Pass in parsed HyperOptArgumentParser to the model
+        :param hparams:
+        """
+        # init superclass
+        super(LightningTestModelBase, self).__init__()
+        self.hparams = hparams
+
+        self.batch_size = hparams.batch_size
+
+        # if you specify an example input, the summary will show input/output for each layer
+        self.example_input_array = torch.rand(5, 28 * 28)
+
+        # remove to test warning for dist sampler
+        self.force_remove_distributed_sampler = force_remove_distributed_sampler
+
+        # 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):
+        """
+        No special modification required for lightning, define as you normally would
+        :param x:
+        :return:
+        """
+
+        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, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop
+        :param batch:
+        :return:
+        """
+        # forward pass
+        x, y = batch
+        x = x.view(x.size(0), -1)
+
+        y_hat = self.forward(x)
+
+        # calculate loss
+        loss_val = self.loss(y, y_hat)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if self.trainer.batch_nb % 1 == 0:
+            output = OrderedDict({
+                'loss': loss_val,
+                'progress_bar': {'some_val': loss_val * loss_val},
+                'log': {'train_some_val': loss_val * loss_val},
+            })
+
+            return output
+        if self.trainer.batch_nb % 2 == 0:
+            return loss_val
+
+    # ---------------------
+    # TRAINING SETUP
+    # ---------------------
+    def configure_optimizers(self):
+        """
+        return whatever optimizers we want here.
+        :return: list of optimizers
+        """
+        # try no scheduler for this model (testing purposes)
+        if self.hparams.optimizer_name == 'lbfgs':
+            optimizer = optim.LBFGS(self.parameters(), lr=self.hparams.learning_rate)
+        else:
+            optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+
+        # test returning only 1 list instead of 2
+        return optimizer
+
+    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)
+
+        # when using multi-node we need to add the datasampler
+        train_sampler = None
+        batch_size = self.hparams.batch_size
+
+        try:
+            if self.use_ddp and not self.force_remove_distributed_sampler:
+                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+                batch_size = batch_size // self.trainer.world_size  # scale batch size
+        except Exception:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
+            dataset=dataset,
+            batch_size=batch_size,
+            shuffle=should_shuffle,
+            sampler=train_sampler
+        )
+
+        return loader
+
+    @data_loader
+    def train_dataloader(self):
+        return self._dataloader(train=True)
+
+    @staticmethod
+    def add_model_specific_args(parent_parser, root_dir):  # pragma: no cover
+        """
+        Parameters you define here will be available to your model through self.hparams
+        :param parent_parser:
+        :param root_dir:
+        :return:
+        """
+        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+
+        # param overwrites
+        # parser.set_defaults(gradient_clip_val=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, type=int)
+        parser.add_argument('--out_features', default=10, type=int)
+        # use 500 for CPU, 50000 for GPU to see speed difference
+        parser.add_argument('--hidden_dim', default=50000, type=int)
+
+        # 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 * 8, type=float,
+                        options=[0.0001, 0.0005, 0.001, 0.005],
+                        tunable=False)
+        parser.opt_list('--optimizer_name', default='adam', type=str,
+                        options=['adam'], tunable=False)
+
+        # if using 2 nodes with 4 gpus each the batch size here
+        #  (256) will be 256 / (2*8) = 16 per gpu
+        parser.opt_list('--batch_size', default=256 * 8, type=int,
+                        options=[32, 64, 128, 256], tunable=False,
+                        help='batch size will be divided over all gpus being used across all nodes')
+        return parser

pytorch_lightning/testing/lm_test_module_mixins.py

@@ -0,0 +1,381 @@
+from collections import OrderedDict
+
+import torch
+
+from pytorch_lightning import data_loader
+
+
+class LightningValidationStepMixin:
+    """
+    Add val_dataloader and validation_step methods for the case
+    when val_dataloader returns a single dataloader
+    """
+
+    @data_loader
+    def val_dataloader(self):
+        return self._dataloader(train=False)
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = 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)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_idx % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return val_acc
+
+        if batch_idx % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
+
+
+class LightningValidationMixin(LightningValidationStepMixin):
+    """
+    Add val_dataloader, validation_step, and validation_end methods for the case
+    when val_dataloader returns a single dataloader
+    """
+
+    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:
+        """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        val_loss_mean = 0
+        val_acc_mean = 0
+        for output in outputs:
+            val_loss = output['val_loss']
+
+            # reduce manually when using dp
+            if self.trainer.use_dp or self.trainer.use_ddp2:
+                val_loss = torch.mean(val_loss)
+            val_loss_mean += val_loss
+
+            # reduce manually when using dp
+            val_acc = output['val_acc']
+            if self.trainer.use_dp or self.trainer.use_ddp2:
+                val_acc = torch.mean(val_acc)
+
+            val_acc_mean += val_acc
+
+        val_loss_mean /= len(outputs)
+        val_acc_mean /= len(outputs)
+
+        tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        results = {'progress_bar': tqdm_dict, 'log': tqdm_dict}
+        return results
+
+
+class LightningValidationStepMultipleDataloadersMixin:
+    """
+    Add val_dataloader and validation_step methods for the case
+    when val_dataloader returns multiple dataloaders
+    """
+
+    @data_loader
+    def val_dataloader(self):
+        return [self._dataloader(train=False), self._dataloader(train=False)]
+
+    def validation_step(self, batch, batch_idx, dataloader_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = 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)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_idx % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return val_acc
+
+        if batch_idx % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
+        if batch_idx % 5 == 0:
+            output = OrderedDict({
+                f'val_loss_{dataloader_idx}': loss_val,
+                f'val_acc_{dataloader_idx}': val_acc,
+            })
+            return output
+
+
+class LightningValidationMultipleDataloadersMixin(LightningValidationStepMultipleDataloadersMixin):
+    """
+    Add val_dataloader, validation_step, and validation_end methods for the case
+    when val_dataloader returns multiple dataloaders
+    """
+
+    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:
+        """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        val_loss_mean = 0
+        val_acc_mean = 0
+        i = 0
+        for dl_output in outputs:
+            for output in dl_output:
+                val_loss = output['val_loss']
+
+                # reduce manually when using dp
+                if self.trainer.use_dp:
+                    val_loss = torch.mean(val_loss)
+                val_loss_mean += val_loss
+
+                # reduce manually when using dp
+                val_acc = output['val_acc']
+                if self.trainer.use_dp:
+                    val_acc = torch.mean(val_acc)
+
+                val_acc_mean += val_acc
+                i += 1
+
+        val_loss_mean /= i
+        val_acc_mean /= i
+
+        tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        result = {'progress_bar': tqdm_dict}
+        return result
+
+
+class LightningTestStepMixin:
+
+    @data_loader
+    def test_dataloader(self):
+        return self._dataloader(train=False)
+
+    def test_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_test = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        test_acc = torch.tensor(test_acc)
+
+        if self.on_gpu:
+            test_acc = test_acc.cuda(loss_test.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_test = loss_test.unsqueeze(0)
+            test_acc = test_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_idx % 1 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return test_acc
+
+        if batch_idx % 3 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+                'test_dic': {'test_loss_a': loss_test}
+            })
+            return output
+
+
+class LightningTestMixin(LightningTestStepMixin):
+    def test_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from test_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        test_loss_mean = 0
+        test_acc_mean = 0
+        for output in outputs:
+            test_loss = output['test_loss']
+
+            # reduce manually when using dp
+            if self.trainer.use_dp:
+                test_loss = torch.mean(test_loss)
+            test_loss_mean += test_loss
+
+            # reduce manually when using dp
+            test_acc = output['test_acc']
+            if self.trainer.use_dp:
+                test_acc = torch.mean(test_acc)
+
+            test_acc_mean += test_acc
+
+        test_loss_mean /= len(outputs)
+        test_acc_mean /= len(outputs)
+
+        tqdm_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+        result = {'progress_bar': tqdm_dict}
+        return result
+
+
+class LightningTestStepMultipleDataloadersMixin:
+
+    @data_loader
+    def test_dataloader(self):
+        return [self._dataloader(train=False), self._dataloader(train=False)]
+
+    def test_step(self, batch, batch_idx, dataloader_idx):
+        """
+        Lightning calls this inside the validation loop
+        :param batch:
+        :return:
+        """
+        x, y = batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_test = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        test_acc = torch.tensor(test_acc)
+
+        if self.on_gpu:
+            test_acc = test_acc.cuda(loss_test.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_test = loss_test.unsqueeze(0)
+            test_acc = test_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if batch_idx % 1 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+            })
+            return output
+        if batch_idx % 2 == 0:
+            return test_acc
+
+        if batch_idx % 3 == 0:
+            output = OrderedDict({
+                'test_loss': loss_test,
+                'test_acc': test_acc,
+                'test_dic': {'test_loss_a': loss_test}
+            })
+            return output
+        if batch_idx % 5 == 0:
+            output = OrderedDict({
+                f'test_loss_{dataloader_idx}': loss_test,
+                f'test_acc_{dataloader_idx}': test_acc,
+            })
+            return output
+
+
+class LightningTestMultipleDataloadersMixin(LightningTestStepMultipleDataloadersMixin):
+    def test_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from test_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+        test_loss_mean = 0
+        test_acc_mean = 0
+        i = 0
+        for dl_output in outputs:
+            for output in dl_output:
+                test_loss = output['test_loss']
+
+                # reduce manually when using dp
+                if self.trainer.use_dp:
+                    test_loss = torch.mean(test_loss)
+                test_loss_mean += test_loss
+
+                # reduce manually when using dp
+                test_acc = output['test_acc']
+                if self.trainer.use_dp:
+                    test_acc = torch.mean(test_acc)
+
+                test_acc_mean += test_acc
+                i += 1
+
+        test_loss_mean /= i
+        test_acc_mean /= i
+
+        tqdm_dict = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()}
+        result = {'progress_bar': tqdm_dict}
+        return result

pytorch_lightning/trainer/amp_mixin.py

@@ -0,0 +1,25 @@
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+import logging
+
+
+class TrainerAMPMixin(object):
+
+    def init_amp(self, use_amp):
+        self.use_amp = use_amp and APEX_AVAILABLE
+        if self.use_amp:
+            logging.info('using 16bit precision')
+
+        if use_amp and not APEX_AVAILABLE:  # pragma: no cover
+            msg = """
+            You set use_amp=True but do not have apex installed.
+            Install apex first using this guide and rerun with use_amp=True:
+            https://github.com/NVIDIA/apex#linux
+
+            this run will NOT use 16 bit precision
+            """
+            raise ModuleNotFoundError(msg)

pytorch_lightning/trainer/callback_config_mixin.py

@@ -0,0 +1,73 @@
+import os
+
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.logging import TestTubeLogger
+
+
+class TrainerCallbackConfigMixin(object):
+    def configure_checkpoint_callback(self):
+        """
+        Weight path set in this priority:
+        Checkpoint_callback's path (if passed in).
+        User provided weights_saved_path
+        Otherwise use os.getcwd()
+        """
+        if self.checkpoint_callback is True:
+            # init a default one
+            if self.logger is not None:
+                ckpt_path = os.path.join(
+                    self.default_save_path,
+                    self.logger.name,
+                    f'version_{self.logger.version}',
+                    "checkpoints"
+                )
+            else:
+                ckpt_path = os.path.join(self.default_save_path, "checkpoints")
+
+            self.checkpoint_callback = ModelCheckpoint(
+                filepath=ckpt_path
+            )
+        elif self.checkpoint_callback is False:
+            self.checkpoint_callback = None
+
+        if self.checkpoint_callback:
+            # set the path for the callbacks
+            self.checkpoint_callback.save_function = self.save_checkpoint
+
+            # if checkpoint callback used, then override the weights path
+            self.weights_save_path = self.checkpoint_callback.filepath
+
+        # if weights_save_path is still none here, set to current working dir
+        if self.weights_save_path is None:
+            self.weights_save_path = self.default_save_path
+
+    def configure_early_stopping(self, early_stop_callback, logger):
+        if early_stop_callback is True:
+            self.early_stop_callback = EarlyStopping(
+                monitor='val_loss',
+                patience=3,
+                verbose=True,
+                mode='min'
+            )
+            self.enable_early_stop = True
+        elif not early_stop_callback:
+            self.early_stop_callback = None
+            self.enable_early_stop = False
+        else:
+            self.early_stop_callback = early_stop_callback
+            self.enable_early_stop = True
+
+        # configure logger
+        if logger is True:
+            # default logger
+            self.logger = TestTubeLogger(
+                save_dir=self.default_save_path,
+                version=self.slurm_job_id,
+                name='lightning_logs'
+            )
+            self.logger.rank = 0
+        elif logger is False:
+            self.logger = None
+        else:
+            self.logger = logger
+            self.logger.rank = 0

pytorch_lightning/trainer/data_loading_mixin.py

@@ -0,0 +1,190 @@
+import warnings
+
+import torch.distributed as dist
+from torch.utils.data import IterableDataset
+from torch.utils.data.distributed import DistributedSampler
+
+from pytorch_lightning.utilities.debugging import MisconfigurationException
+
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+
+
+class TrainerDataLoadingMixin(object):
+    def init_train_dataloader(self, model):
+        """
+        Dataloaders are provided by the model
+        :param model:
+        :return:
+        """
+        self.get_train_dataloader = model.train_dataloader
+
+        # determine number of training batches
+        if isinstance(self.get_train_dataloader(), IterableDataset):
+            self.nb_training_batches = float('inf')
+        else:
+            self.nb_training_batches = len(self.get_train_dataloader())
+            self.nb_training_batches = int(self.nb_training_batches * self.train_percent_check)
+
+        # determine when to check validation
+        # if int passed in, val checks that often
+        # otherwise, it checks in [0, 1.0] % range of a training epoch
+        if isinstance(self.val_check_interval, int):
+            self.val_check_batch = self.val_check_interval
+        else:
+            self.val_check_batch = int(self.nb_training_batches * self.val_check_interval)
+            self.val_check_batch = max(1, self.val_check_batch)
+
+        on_ddp = self.use_ddp or self.use_ddp2
+        if on_ddp and not isinstance(self.get_train_dataloader().sampler, DistributedSampler):
+            msg = """
+            You're using multiple gpus and multiple nodes without using a DistributedSampler
+            to assign a subset of your data to each process. To silence this warning, pass a
+            DistributedSampler to your DataLoader.
+
+            ie: this:
+            dataset = myDataset()
+            dataloader = Dataloader(dataset)
+
+            becomes:
+            dataset = myDataset()
+            dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+            dataloader = Dataloader(dataset, sampler=dist_sampler)
+
+            If you want each process to load the full dataset, ignore this warning.
+            """
+            if msg not in self.shown_warnings and self.proc_rank == 0:
+                self.shown_warnings.add(msg)
+                warnings.warn(msg)
+
+    def init_val_dataloader(self, model):
+        """
+        Dataloaders are provided by the model
+        :param model:
+        :return:
+        """
+        self.get_val_dataloaders = model.val_dataloader
+
+        # determine number of validation batches
+        # val datasets could be none, 1 or 2+
+        if self.get_val_dataloaders() is not None:
+            self.nb_val_batches = sum(len(dataloader) for dataloader in self.get_val_dataloaders())
+            self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
+            self.nb_val_batches = max(1, self.nb_val_batches)
+
+        on_ddp = self.use_ddp or self.use_ddp2
+        if on_ddp and self.get_val_dataloaders() is not None:
+            for dataloader in self.get_val_dataloaders():
+                if not isinstance(dataloader.sampler, DistributedSampler):
+                    msg = """
+                    Your val_dataloader(s) don't use DistributedSampler.
+
+                    You're using multiple gpus and multiple nodes without using a
+                    DistributedSampler to assign a subset of your data to each process.
+                    To silence this warning, pass a DistributedSampler to your DataLoader.
+
+                    ie: this:
+                    dataset = myDataset()
+                    dataloader = Dataloader(dataset)
+
+                    becomes:
+                    dataset = myDataset()
+                    dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+                    dataloader = Dataloader(dataset, sampler=dist_sampler)
+
+                    If you want each process to load the full dataset, ignore this warning.
+                    """
+                    if msg not in self.shown_warnings and self.proc_rank == 0:
+                        self.shown_warnings.add(msg)
+                        warnings.warn(msg)
+                    break
+
+    def init_test_dataloader(self, model):
+        """
+        Dataloaders are provided by the model
+        :param model:
+        :return:
+        """
+
+        self.get_test_dataloaders = model.test_dataloader
+
+        # determine number of test batches
+        if self.get_test_dataloaders() is not None:
+            len_sum = sum(len(dataloader) for dataloader in self.get_test_dataloaders())
+            self.nb_test_batches = len_sum
+            self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
+            self.nb_test_batches = max(1, self.nb_test_batches)
+
+        on_ddp = self.use_ddp or self.use_ddp2
+        if on_ddp and self.get_test_dataloaders() is not None:
+            for dataloader in self.get_test_dataloaders():
+                if not isinstance(dataloader.sampler, DistributedSampler):
+                    msg = """
+                    Your test_dataloader(s) don't use DistributedSampler.
+
+                    You're using multiple gpus and multiple nodes without using a
+                    DistributedSampler to assign a subset of your data to each process.
+                    To silence this warning, pass a DistributedSampler to your DataLoader.
+
+                    ie: this:
+                    dataset = myDataset()
+                    dataloader = Dataloader(dataset)
+
+                    becomes:
+                    dataset = myDataset()
+                    dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+                    dataloader = Dataloader(dataset, sampler=dist_sampler)
+
+                    If you want each process to load the full dataset, ignore this warning.
+                    """
+                    if msg not in self.shown_warnings and self.proc_rank == 0:
+                        self.shown_warnings.add(msg)
+                        warnings.warn(msg)
+                    break
+
+    def get_dataloaders(self, model):
+        """
+        Dataloaders are provided by the model
+        :param model:
+        :return:
+        """
+
+        self.init_train_dataloader(model)
+        self.init_test_dataloader(model)
+        self.init_val_dataloader(model)
+
+        if self.use_ddp or self.use_ddp2:
+            # wait for all processes to catch up
+            dist.barrier()
+
+            # load each dataloader
+            self.get_train_dataloader()
+            self.get_test_dataloaders()
+            self.get_val_dataloaders()
+
+        # support IterableDataset for train data
+        self.is_iterable_train_dataloader = isinstance(self.get_train_dataloader(), IterableDataset)
+        if self.is_iterable_train_dataloader and not isinstance(self.val_check_interval, int):
+            m = '''
+            When using an iterableDataset for train_dataloader,
+            Trainer(val_check_interval) must be an int.
+            An int k specifies checking validation every k training batches
+            '''
+            raise MisconfigurationException(m)
+
+    def determine_data_use_amount(self, train_percent_check, val_percent_check,
+                                  test_percent_check, overfit_pct):
+        """
+        Use less data for debugging purposes
+        """
+        self.train_percent_check = train_percent_check
+        self.val_percent_check = val_percent_check
+        self.test_percent_check = test_percent_check
+        if overfit_pct > 0:
+            self.train_percent_check = overfit_pct
+            self.val_percent_check = overfit_pct
+            self.test_percent_check = overfit_pct

pytorch_lightning/trainer/ddp_mixin.py

@@ -0,0 +1,193 @@
+import os
+import re
+import warnings
+import logging
+
+import torch
+
+from pytorch_lightning.utilities.debugging import MisconfigurationException
+
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+
+
+class TrainerDDPMixin(object):
+    def set_distributed_mode(self, distributed_backend, nb_gpu_nodes):
+        # skip for CPU
+        if self.num_gpus == 0:
+            return
+
+        # single GPU case
+        # in single gpu case we allow ddp so we can train on multiple
+        # nodes, 1 gpu per node
+        if self.num_gpus == 1:
+            self.single_gpu = True
+
+            if distributed_backend is not None:
+                self.use_dp = distributed_backend == 'dp'
+                self.use_ddp = distributed_backend == 'ddp'
+                self.use_ddp2 = distributed_backend == 'ddp2'
+
+                # disable single gpu when using ddp2
+                if self.use_ddp2:
+                    self.single_gpu = False
+
+        # multiple GPU case
+        elif self.num_gpus > 1:
+            if distributed_backend is not None:
+                # DP, DDP case
+                self.use_dp = distributed_backend == 'dp'
+                self.use_ddp = distributed_backend == 'ddp'
+                self.use_ddp2 = distributed_backend == 'ddp2'
+
+            elif distributed_backend is None:
+                m = 'When using multiple GPUs set ' \
+                    'Trainer(distributed_backend=dp) (or ddp)'
+                raise MisconfigurationException(m)
+
+        # throw error to force user ddp or ddp2 choice
+        if nb_gpu_nodes > 1 and not (self.use_ddp2 or self.use_ddp):  # pragma: no cover
+            w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
+                'Switching to DistributedDataParallel for you. ' \
+                'To silence this warning set distributed_backend=ddp' \
+                'or distributed_backend=ddp2'
+            raise MisconfigurationException(w)
+
+        logging.info(f'gpu available: {torch.cuda.is_available()}, used: {self.on_gpu}')
+
+    def configure_slurm_ddp(self, nb_gpu_nodes):
+        self.is_slurm_managing_tasks = False
+
+        # extract SLURM flag vars
+        # whenever we have the correct number of tasks, we let slurm manage processes
+        # otherwise we launch the required number of processes
+        if self.use_ddp:
+            self.nb_requested_gpus = self.num_gpus * nb_gpu_nodes
+            self.nb_slurm_tasks = 0
+            try:
+                self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+                self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
+
+                # in interactive mode we don't manage tasks
+                job_name = os.environ['SLURM_JOB_NAME']
+                if job_name == 'bash':
+                    self.is_slurm_managing_tasks = False
+
+            except Exception:
+                # likely not on slurm, so set the slurm managed flag to false
+                self.is_slurm_managing_tasks = False
+
+        # used for tests only, set this flag to simulate slurm managing a task
+        try:
+            should_fake = int(os.environ['FAKE_SLURM_MANAGING_TASKS'])
+            if should_fake:
+                self.is_slurm_managing_tasks = True
+        except Exception as e:
+            pass
+
+    def set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids):
+        if data_parallel_device_ids is None:
+            return
+
+        # set the correct cuda visible devices (using pci order)
+        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+        # when slurm is managing the task it sets the visible devices
+        if not is_slurm_managing_tasks:
+            if type(data_parallel_device_ids) is int:
+                id_str = ','.join(str(x) for x in list(range(data_parallel_device_ids)))
+                os.environ["CUDA_VISIBLE_DEVICES"] = id_str
+            else:
+                gpu_str = ','.join([str(x) for x in data_parallel_device_ids])
+                os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str
+
+        logging.info(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
+
+    def ddp_train(self, gpu_nb, model):
+        """
+        Entry point into a DP thread
+        :param gpu_nb:
+        :param model:
+        :param cluster_obj:
+        :return:
+        """
+        # node rank using relative slurm id
+        # otherwise default to node rank 0
+        try:
+            node_id = os.environ['SLURM_NODEID']
+            self.node_rank = int(node_id)
+        except Exception:
+            self.node_rank = 0
+
+        # show progressbar only on progress_rank 0
+        self.show_progress_bar = self.show_progress_bar and self.node_rank == 0 and gpu_nb == 0
+
+        # determine which process we are and world size
+        if self.use_ddp:
+            self.proc_rank = self.node_rank * self.num_gpus + gpu_nb
+            self.world_size = self.nb_gpu_nodes * self.num_gpus
+
+        elif self.use_ddp2:
+            self.proc_rank = self.node_rank
+            self.world_size = self.nb_gpu_nodes
+
+        # let the exp know the rank to avoid overwriting logs
+        if self.logger is not None:
+            self.logger.rank = self.proc_rank
+
+        # set up server using proc 0's ip address
+        # try to init for 20 times at max in case ports are taken
+        # where to store ip_table
+        model.trainer = self
+        model.init_ddp_connection(self.proc_rank, self.world_size)
+
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+
+        # MODEL
+        # copy model to each gpu
+        if self.distributed_backend == 'ddp':
+            torch.cuda.set_device(gpu_nb)
+        model.cuda(gpu_nb)
+
+        # set model properties before going into wrapper
+        self.copy_trainer_model_properties(model)
+
+        # override root GPU
+        self.root_gpu = gpu_nb
+
+        # AMP
+        # run through amp wrapper before going to distributed DP
+        if self.use_amp:
+            # An example
+            model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level)
+            self.optimizers = optimizers
+
+        # DDP2 uses all GPUs on the machine
+        if self.distributed_backend == 'ddp':
+            device_ids = [gpu_nb]
+        elif self.use_ddp2:
+            device_ids = self.data_parallel_device_ids
+
+        # allow user to configure ddp
+        model = model.configure_ddp(model, device_ids)
+
+        # continue training routine
+        self.run_pretrain_routine(model)
+
+    def resolve_root_node_address(self, root_node):
+        if '[' in root_node:
+            name = root_node.split('[')[0]
+            number = root_node.split(',')[0]
+            if '-' in number:
+                number = number.split('-')[0]
+
+            number = re.sub('[^0-9]', '', number)
+            root_node = name + number
+
+        return root_node

pytorch_lightning/trainer/dp_mixin.py

@@ -0,0 +1,217 @@
+import torch
+
+from pytorch_lightning.pt_overrides.override_data_parallel import (
+    LightningDistributedDataParallel, LightningDataParallel)
+from pytorch_lightning.utilities.debugging import MisconfigurationException
+
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+
+
+class TrainerDPMixin(object):
+    def copy_trainer_model_properties(self, model):
+        if isinstance(model, LightningDataParallel):
+            ref_model = model.module
+        elif isinstance(model, LightningDistributedDataParallel):
+            ref_model = model.module
+        else:
+            ref_model = model
+
+        for m in [model, ref_model]:
+            m.trainer = self
+            m.on_gpu = self.on_gpu
+            m.use_dp = self.use_dp
+            m.use_ddp2 = self.use_ddp2
+            m.use_ddp = self.use_ddp
+            m.use_amp = self.use_amp
+            m.testing = self.testing
+            m.single_gpu = self.single_gpu
+
+    def transfer_batch_to_gpu(self, batch, gpu_id):
+        # base case: object can be directly moved using `cuda` or `to`
+        if callable(getattr(batch, 'cuda', None)):
+            return batch.cuda(gpu_id)
+
+        elif callable(getattr(batch, 'to', None)):
+            return batch.to(torch.device('cuda', gpu_id))
+
+        # when list
+        elif isinstance(batch, list):
+            for i, x in enumerate(batch):
+                batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
+            return batch
+
+        # when tuple
+        elif isinstance(batch, tuple):
+            batch = list(batch)
+            for i, x in enumerate(batch):
+                batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
+            return tuple(batch)
+
+        # when dict
+        elif isinstance(batch, dict):
+            for k, v in batch.items():
+                batch[k] = self.transfer_batch_to_gpu(v, gpu_id)
+
+            return batch
+
+        # nothing matches, return the value as is without transform
+        return batch
+
+    def single_gpu_train(self, model):
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+
+        model.cuda(self.root_gpu)
+
+        if self.use_amp:
+            # An example
+            model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level)
+            self.optimizers = optimizers
+
+        self.run_pretrain_routine(model)
+
+    def dp_train(self, model):
+
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+
+        model.cuda(self.root_gpu)
+
+        # check for this bug (amp + dp + !01 doesn't work)
+        # https://github.com/NVIDIA/apex/issues/227
+        if self.use_dp and self.use_amp:
+            m = f"""
+            Amp level {self.amp_level} with DataParallel is not supported.
+            See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227.
+            We recommend you switch to ddp if you want to use amp
+            """
+            raise MisconfigurationException(m)
+
+        # create list of device ids
+        device_ids = self.data_parallel_device_ids
+        if type(device_ids) is int:
+            device_ids = list(range(device_ids))
+
+        model = LightningDataParallel(model, device_ids=device_ids)
+
+        self.run_pretrain_routine(model)
+
+
+def normalize_parse_gpu_string_input(s):
+    if type(s) is str:
+        if s == '-1':
+            return -1
+        else:
+            return [int(x.strip()) for x in s.split(',')]
+    else:
+        return s
+
+
+def get_all_available_gpus():
+    """
+    :return: a list of all available gpus
+    """
+    return list(range(torch.cuda.device_count()))
+
+
+def check_gpus_data_type(gpus):
+    """
+    :param gpus: gpus parameter as passed to the Trainer
+        Function checks that it is one of: None, Int, String or List
+        Throws otherwise
+    :return: return unmodified gpus variable
+    """
+
+    if (gpus is not None and
+        type(gpus) is not int and
+        type(gpus) is not str and
+        type(gpus) is not list):    # noqa E129
+        raise MisconfigurationException("GPUs must be int, string or list of ints or None.")
+
+
+def normalize_parse_gpu_input_to_list(gpus):
+    assert gpus is not None
+    if isinstance(gpus, list):
+        return gpus
+    else:  # must be an int
+        if not gpus:  # gpus==0
+            return None
+        elif gpus == -1:
+            return get_all_available_gpus()
+        else:
+            return list(range(gpus))
+
+
+def sanitize_gpu_ids(gpus):
+    """
+    :param gpus: list of ints corresponding to GPU indices
+        Checks that each of the GPUs in the list is actually available.
+        Throws if any of the GPUs is not available.
+    :return: unmodified gpus variable
+    """
+    all_available_gpus = get_all_available_gpus()
+    for gpu in gpus:
+        if gpu not in all_available_gpus:
+            message = f"""
+            Non-available gpu index {gpu} specified:
+            Available gpu indices are: {all_available_gpus}
+            """
+            raise MisconfigurationException(message)
+    return gpus
+
+
+def parse_gpu_ids(gpus):
+    """
+    :param gpus: Int, string or list
+        An int -1 or string '-1' indicate that all available GPUs should be used.
+        A list of ints or a string containing list of comma separated integers
+        indicates specific GPUs to use
+        An int 0 means that no GPUs should be used
+        Any int N > 0 indicates that GPUs [0..N) should be used.
+    :return: List of gpus to be used
+
+        If no GPUs are available but the value of gpus variable indicates request for GPUs
+        then a misconfiguration exception is raised.
+    """
+
+    # Check that gpus param is None, Int, String or List
+    check_gpus_data_type(gpus)
+
+    # Handle the case when no gpus are requested
+    if gpus is None or type(gpus) is int and gpus == 0:
+        return None
+
+    # We know user requested GPUs therefore if some of the
+    # requested GPUs are not available an exception is thrown.
+
+    gpus = normalize_parse_gpu_string_input(gpus)
+    gpus = normalize_parse_gpu_input_to_list(gpus)
+    gpus = sanitize_gpu_ids(gpus)
+
+    if not gpus:
+        raise MisconfigurationException("GPUs requested but non are available.")
+    return gpus
+
+
+def determine_root_gpu_device(gpus):
+    """
+    :param gpus: non empty list of ints representing which gpus to use
+    :return: designated root GPU device
+    """
+    if gpus is None:
+        return None
+
+    assert isinstance(gpus, list), "gpus should be a list"
+    assert len(gpus), "gpus should be a non empty list"
+
+    # set root gpu
+    root_gpu = gpus[0]
+
+    return root_gpu

pytorch_lightning/trainer/evaluation_loop_mixin.py

@@ -0,0 +1,192 @@
+import torch
+import tqdm
+
+from pytorch_lightning.utilities.debugging import MisconfigurationException
+
+
+class TrainerEvaluationLoopMixin(object):
+
+    def evaluate(self, model, dataloaders, max_batches, test=False):
+        """
+        Run evaluation code
+        :param model: PT model
+        :param dataloaders: list of PT dataloaders
+        :param max_batches: Scalar
+        :param test: boolean
+        :return:
+        """
+        # enable eval mode
+        model.zero_grad()
+        model.eval()
+
+        # copy properties for forward overrides
+        self.copy_trainer_model_properties(model)
+
+        # disable gradients to save memory
+        torch.set_grad_enabled(False)
+
+        # bookkeeping
+        outputs = []
+
+        # run training
+        for dataloader_idx, dataloader in enumerate(dataloaders):
+            dl_outputs = []
+            for batch_idx, batch in enumerate(dataloader):
+
+                if batch is None:  # pragma: no cover
+                    continue
+
+                # stop short when on fast_dev_run (sets max_batch=1)
+                if batch_idx >= max_batches:
+                    break
+
+                # -----------------
+                # RUN EVALUATION STEP
+                # -----------------
+                output = self.evaluation_forward(model,
+                                                 batch,
+                                                 batch_idx,
+                                                 dataloader_idx,
+                                                 test)
+
+                # track outputs for collation
+                dl_outputs.append(output)
+
+                # batch done
+                if test:
+                    self.test_progress_bar.update(1)
+                else:
+                    self.val_progress_bar.update(1)
+                    self.main_progress_bar.update(1)
+            outputs.append(dl_outputs)
+
+        eval_results = {}
+
+        # with a single dataloader don't pass an array
+        if len(dataloaders) == 1:
+            outputs = outputs[0]
+
+        # give model a chance to do something with the outputs (and method defined)
+        model = self.get_model()
+        if test and self.is_overriden('test_end'):
+            eval_results = model.test_end(outputs)
+        elif self.is_overriden('validation_end'):
+            eval_results = model.validation_end(outputs)
+
+        # enable train mode again
+        model.train()
+
+        # enable gradients to save memory
+        torch.set_grad_enabled(True)
+
+        return eval_results
+
+    def run_evaluation(self, test=False):
+        # when testing make sure user defined a test step
+        can_run_test_step = False
+        if test:
+            can_run_test_step = self.is_overriden('test_step') and self.is_overriden('test_end')
+            if not can_run_test_step:
+                m = '''You called .test() without defining a test step or test_end.
+                Please define and try again'''
+                raise MisconfigurationException(m)
+
+        # validate only if model has validation_step defined
+        # test only if test_step or validation_step are defined
+        run_val_step = self.is_overriden('validation_step')
+
+        if run_val_step or can_run_test_step:
+
+            # hook
+            model = self.get_model()
+            model.on_pre_performance_check()
+
+            # select dataloaders
+            if test:
+                dataloaders = self.get_test_dataloaders()
+                max_batches = self.nb_test_batches
+            else:
+                # val
+                dataloaders = self.get_val_dataloaders()
+                max_batches = self.nb_val_batches
+
+            # cap max batches to 1 when using fast_dev_run
+            if self.fast_dev_run:
+                max_batches = 1
+
+            # init validation or test progress bar
+            # main progress bar will already be closed when testing so initial position is free
+            position = 2 * self.process_position + (not test)
+            desc = 'Testing' if test else 'Validating'
+            pbar = tqdm.tqdm(desc=desc, total=max_batches, leave=test, position=position,
+                             disable=not self.show_progress_bar, dynamic_ncols=True,
+                             unit='batch')
+            setattr(self, f'{"test" if test else "val"}_progress_bar', pbar)
+
+            # run evaluation
+            eval_results = self.evaluate(self.model,
+                                         dataloaders,
+                                         max_batches,
+                                         test)
+            _, prog_bar_metrics, log_metrics, callback_metrics, _ = self.process_output(
+                eval_results)
+
+            # add metrics to prog bar
+            self.add_tqdm_metrics(prog_bar_metrics)
+
+            # log metrics
+            self.log_metrics(log_metrics, {})
+
+            # track metrics for callbacks
+            self.callback_metrics = callback_metrics
+
+            # hook
+            model.on_post_performance_check()
+
+            # add model specific metrics
+            tqdm_metrics = self.training_tqdm_dict
+            if not test:
+                self.main_progress_bar.set_postfix(**tqdm_metrics)
+
+            # close progress bar
+            if test:
+                self.test_progress_bar.close()
+            else:
+                self.val_progress_bar.close()
+
+        # model checkpointing
+        if self.proc_rank == 0 and self.checkpoint_callback is not None and not test:
+            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch,
+                                                  logs=self.callback_metrics)
+
+    def evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test=False):
+        # make dataloader_idx arg in validation_step optional
+        args = [batch, batch_idx]
+
+        if test and len(self.get_test_dataloaders()) > 1:
+            args.append(dataloader_idx)
+
+        elif not test and len(self.get_val_dataloaders()) > 1:
+            args.append(dataloader_idx)
+
+        # handle DP, DDP forward
+        if self.use_ddp or self.use_dp or self.use_ddp2:
+            output = model(*args)
+            return output
+
+        # single GPU
+        if self.single_gpu:
+            # for single GPU put inputs on gpu manually
+            root_gpu = 0
+            if type(self.data_parallel_device_ids) is list:
+                root_gpu = self.data_parallel_device_ids[0]
+            batch = self.transfer_batch_to_gpu(batch, root_gpu)
+            args[0] = batch
+
+        # CPU
+        if test:
+            output = model.test_step(*args)
+        else:
+            output = model.validation_step(*args)
+
+        return output

pytorch_lightning/trainer/ignored_warnings.py

@@ -0,0 +1,14 @@
+import warnings
+
+
+def ignore_scalar_return_in_dp():
+    # Users get confused by this warning so we silence it
+    m_1 = """
+    Was asked to gather along dimension 0, but all
+    input tensors were scalars; will instead unsqueeze
+    and return a vector.
+    """
+    warnings.filterwarnings('ignore', message=m_1)
+
+
+ignore_scalar_return_in_dp()

pytorch_lightning/trainer/logging_mixin.py

@@ -0,0 +1,167 @@
+import torch
+
+from pytorch_lightning.root_module import memory
+
+
+class TrainerLoggingMixin(object):
+
+    def log_metrics(self, metrics, grad_norm_dic):
+        """
+        Logs the metric dict passed in
+        :param metrics:
+        :param grad_norm_dic:
+        :return:
+        """
+        # added metrics by Lightning for convenience
+        metrics['epoch'] = self.current_epoch
+
+        # add gpu memory
+        if self.on_gpu and self.log_gpu_memory:
+            mem_map = memory.get_memory_profile(self.log_gpu_memory)
+            metrics.update(mem_map)
+
+        # add norms
+        metrics.update(grad_norm_dic)
+
+        # turn all tensors to scalars
+        scalar_metrics = self.metrics_to_scalars(metrics)
+
+        # log actual metrics
+        if self.proc_rank == 0 and self.logger is not None:
+            self.logger.log_metrics(scalar_metrics, step_num=self.global_step)
+            self.logger.save()
+
+    def add_tqdm_metrics(self, metrics):
+        for k, v in metrics.items():
+            if type(v) is torch.Tensor:
+                v = v.item()
+
+            self.tqdm_metrics[k] = v
+
+    def metrics_to_scalars(self, metrics):
+        new_metrics = {}
+        for k, v in metrics.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+
+            if type(v) is dict:
+                v = self.metrics_to_scalars(v)
+
+            new_metrics[k] = v
+
+        return new_metrics
+
+    def process_output(self, output, train=False):
+        """
+        Reduces output according to the training mode.
+        Separates loss from logging and tqdm metrics
+        :param output:
+        :return:
+        """
+        # ---------------
+        # EXTRACT CALLBACK KEYS
+        # ---------------
+        # all keys not progress_bar or log are candidates for callbacks
+        callback_metrics = {}
+        for k, v in output.items():
+            if k not in ['progress_bar', 'log', 'hiddens']:
+                callback_metrics[k] = v
+
+        if train and (self.use_dp or self.use_ddp2):
+            nb_gpus = self.num_gpus
+            callback_metrics = self.reduce_distributed_output(callback_metrics, nb_gpus)
+
+        for k, v in callback_metrics.items():
+            callback_metrics[k] = v.item()
+
+        # ---------------
+        # EXTRACT PROGRESS BAR KEYS
+        # ---------------
+        try:
+            progress_output = output['progress_bar']
+
+            # reduce progress metrics for tqdm when using dp
+            if train and (self.use_dp or self.use_ddp2):
+                nb_gpus = self.num_gpus
+                progress_output = self.reduce_distributed_output(progress_output, nb_gpus)
+
+            progress_bar_metrics = progress_output
+        except Exception:
+            progress_bar_metrics = {}
+
+        # ---------------
+        # EXTRACT LOGGING KEYS
+        # ---------------
+        # extract metrics to log to experiment
+        try:
+            log_output = output['log']
+
+            # reduce progress metrics for tqdm when using dp
+            if train and (self.use_dp or self.use_ddp2):
+                nb_gpus = self.num_gpus
+                log_output = self.reduce_distributed_output(log_output, nb_gpus)
+
+            log_metrics = log_output
+        except Exception:
+            log_metrics = {}
+
+        # ---------------
+        # EXTRACT LOSS
+        # ---------------
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        loss = None
+        if train:
+            try:
+                loss = output['loss']
+            except Exception:
+                if type(output) is torch.Tensor:
+                    loss = output
+                else:
+                    raise RuntimeError(
+                        'No `loss` value in the dictionary returned from `model.training_step()`.'
+                    )
+
+            # when using dp need to reduce the loss
+            if self.use_dp or self.use_ddp2:
+                loss = self.reduce_distributed_output(loss, self.num_gpus)
+
+        # ---------------
+        # EXTRACT HIDDEN
+        # ---------------
+        hiddens = output.get('hiddens')
+
+        # use every metric passed in as a candidate for callback
+        callback_metrics.update(progress_bar_metrics)
+        callback_metrics.update(log_metrics)
+
+        # convert tensors to numpy
+        for k, v in callback_metrics.items():
+            if isinstance(v, torch.Tensor):
+                callback_metrics[k] = v.item()
+
+        return loss, progress_bar_metrics, log_metrics, callback_metrics, hiddens
+
+    def reduce_distributed_output(self, output, nb_gpus):
+        if nb_gpus <= 1:
+            return output
+
+        # when using DP, we get one output per gpu
+        # average outputs and return
+        if type(output) is torch.Tensor:
+            return output.mean()
+
+        for k, v in output.items():
+            # recurse on nested dics
+            if isinstance(output[k], dict):
+                output[k] = self.reduce_distributed_output(output[k], nb_gpus)
+
+            # do nothing when there's a scalar
+            elif isinstance(output[k], torch.Tensor) and output[k].dim() == 0:
+                pass
+
+            # 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

pytorch_lightning/trainer/model_hooks_mixin.py

@@ -0,0 +1,17 @@
+from pytorch_lightning.root_module.root_module import LightningModule
+
+
+class TrainerModelHooksMixin(object):
+
+    def is_function_implemented(self, f_name):
+        model = self.get_model()
+        f_op = getattr(model, f_name, None)
+        return callable(f_op)
+
+    def is_overriden(self, f_name):
+        model = self.get_model()
+        super_object = LightningModule
+
+        # when code pointers are different, it was overriden
+        is_overriden = getattr(model, f_name).__code__ is not getattr(super_object, f_name).__code__
+        return is_overriden

pytorch_lightning/trainer/train_loop_mixin.py

@@ -0,0 +1,309 @@
+import numpy as np
+import tqdm
+
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+
+
+class TrainerTrainLoopMixin(object):
+
+    def train(self):
+        # run all epochs
+        for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
+            # set seed for distributed sampler (enables shuffling for each epoch)
+            if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):
+                self.get_train_dataloader().sampler.set_epoch(epoch_nb)
+
+            # get model
+            model = self.get_model()
+
+            # update training progress in trainer and model
+            model.current_epoch = epoch_nb
+            self.current_epoch = epoch_nb
+
+            # val can be checked multiple times in epoch
+            is_val_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
+            val_checks_per_epoch = self.nb_training_batches // self.val_check_batch
+            val_checks_per_epoch = val_checks_per_epoch if is_val_epoch else 0
+
+            # total batches includes multiple val checks
+            self.total_batches = (self.nb_training_batches +
+                                  self.nb_val_batches * val_checks_per_epoch)
+            self.batch_loss_value = 0  # accumulated grads
+
+            if self.fast_dev_run:
+                # limit the number of batches to 2 (1 train and 1 val) in fast_dev_run
+                nb_iterations = 2
+            elif self.is_iterable_train_dataloader:
+                # for iterable train loader, the progress bar never ends
+                nb_iterations = None
+            else:
+                nb_iterations = self.total_batches
+
+            # reset progress bar
+            # .reset() doesn't work on disabled progress bar so we should check
+            if not self.main_progress_bar.disable:
+                self.main_progress_bar.reset(nb_iterations)
+            desc = f'Epoch {epoch_nb + 1}' if not self.is_iterable_train_dataloader else ''
+            self.main_progress_bar.set_description(desc)
+
+            # changing gradient according accumulation_scheduler
+            self.accumulation_scheduler.on_epoch_begin(epoch_nb, self)
+
+            # -----------------
+            # RUN TNG EPOCH
+            # -----------------
+            self.run_training_epoch()
+
+            # update LR schedulers
+            if self.lr_schedulers is not None:
+                for lr_scheduler in self.lr_schedulers:
+                    lr_scheduler.step(self.current_epoch)
+
+            # early stopping
+            met_min_epochs = epoch_nb > self.min_nb_epochs
+            if self.enable_early_stop and (met_min_epochs or self.fast_dev_run):
+                should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb,
+                                                                    logs=self.callback_metrics)
+                # stop training
+                stop = should_stop and met_min_epochs
+                if stop:
+                    self.main_progress_bar.close()
+                    return
+
+        self.main_progress_bar.close()
+
+        if self.logger is not None:
+            self.logger.finalize("success")
+
+    def run_training_epoch(self):
+        # before epoch hook
+        if self.is_function_implemented('on_epoch_start'):
+            model = self.get_model()
+            model.on_epoch_start()
+
+        # run epoch
+        for batch_nb, batch in enumerate(self.get_train_dataloader()):
+            self.batch_nb = batch_nb
+
+            model = self.get_model()
+            model.global_step = self.global_step
+
+            # ---------------
+            # RUN TRAIN STEP
+            # ---------------
+            output = self.run_training_batch(batch, batch_nb)
+            batch_result, grad_norm_dic, batch_step_metrics = output
+
+            # when returning -1 from train_step, we end epoch early
+            early_stop_epoch = batch_result == -1
+
+            # ---------------
+            # RUN VAL STEP
+            # ---------------
+            is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
+            can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
+            should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch)
+
+            # fast_dev_run always forces val checking after train batch
+            if self.fast_dev_run or should_check_val:
+                self.run_evaluation(test=self.testing)
+
+            # when logs should be saved
+            should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch
+            if should_save_log or self.fast_dev_run:
+                if self.proc_rank == 0 and self.logger is not None:
+                    self.logger.save()
+
+            # when metrics should be logged
+            should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch
+            if should_log_metrics or self.fast_dev_run:
+                # logs user requested information to logger
+                self.log_metrics(batch_step_metrics, grad_norm_dic)
+
+            self.global_step += 1
+            self.total_batch_nb += 1
+
+            # end epoch early
+            # stop when the flag is changed or we've gone past the amount
+            # requested in the batches
+            if early_stop_epoch or self.fast_dev_run:
+                break
+
+            # stop epoch if we limited nb batches
+            met_batch_limit = batch_nb >= self.nb_training_batches
+            if met_batch_limit:
+                break
+
+        # epoch end hook
+        if self.is_function_implemented('on_epoch_end'):
+            model = self.get_model()
+            model.on_epoch_end()
+
+    def run_training_batch(self, batch, batch_nb):
+        # track grad norms
+        grad_norm_dic = {}
+
+        # track all metrics for callbacks
+        all_callback_metrics = []
+
+        # track metrics to log
+        all_log_metrics = []
+
+        if batch is None:
+            return 0, grad_norm_dic
+
+        # hook
+        if self.is_function_implemented('on_batch_start'):
+            model_ref = self.get_model()
+            response = model_ref.on_batch_start(batch)
+
+            if response == -1:
+                return -1, grad_norm_dic
+
+        splits = [batch]
+        if self.truncated_bptt_steps is not None:
+            model_ref = self.get_model()
+            splits = model_ref.tbptt_split_batch(batch, self.truncated_bptt_steps)
+
+        self.hiddens = None
+        for split_nb, split_batch in enumerate(splits):
+            self.split_nb = split_nb
+
+            # call training_step once per optimizer
+            for opt_idx, optimizer in enumerate(self.optimizers):
+
+                # wrap the forward step in a closure so second order methods work
+                def optimizer_closure():
+                    # forward pass
+                    output = self.training_forward(
+                        split_batch, batch_nb, opt_idx, self.hiddens)
+
+                    closure_loss = output[0]
+                    progress_bar_metrics = output[1]
+                    log_metrics = output[2]
+                    callback_metrics = output[3]
+                    self.hiddens = output[4]
+
+                    # accumulate loss
+                    # (if accumulate_grad_batches = 1 no effect)
+                    closure_loss = closure_loss / self.accumulate_grad_batches
+
+                    # backward pass
+                    model_ref = self.get_model()
+                    model_ref.backward(self.use_amp, closure_loss, optimizer)
+
+                    # track metrics for callbacks
+                    all_callback_metrics.append(callback_metrics)
+
+                    # track progress bar metrics
+                    self.add_tqdm_metrics(progress_bar_metrics)
+                    all_log_metrics.append(log_metrics)
+
+                    # insert after step hook
+                    if self.is_function_implemented('on_after_backward'):
+                        model_ref = self.get_model()
+                        model_ref.on_after_backward()
+
+                    return closure_loss
+
+                # calculate loss
+                loss = optimizer_closure()
+
+                # nan grads
+                if self.print_nan_grads:
+                    self.print_nan_gradients()
+
+                # track total loss for logging (avoid mem leaks)
+                self.batch_loss_value += loss.item()
+
+                # gradient update with accumulated gradients
+                if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
+
+                    # track gradient norms when requested
+                    if batch_nb % self.row_log_interval == 0:
+                        if self.track_grad_norm > 0:
+                            model = self.get_model()
+                            grad_norm_dic = model.grad_norm(
+                                self.track_grad_norm)
+
+                    # clip gradients
+                    self.clip_gradients()
+
+                    # calls .step(), .zero_grad()
+                    # override function to modify this behavior
+                    model = self.get_model()
+                    model.optimizer_step(self.current_epoch, batch_nb,
+                                         optimizer, opt_idx, optimizer_closure)
+
+                    # calculate running loss for display
+                    self.running_loss.append(self.batch_loss_value)
+                    self.batch_loss_value = 0
+                    self.avg_loss = np.mean(self.running_loss[-100:])
+
+        # activate batch end hook
+        if self.is_function_implemented('on_batch_end'):
+            model = self.get_model()
+            model.on_batch_end()
+
+        # update progress bar
+        self.main_progress_bar.update(1)
+        self.main_progress_bar.set_postfix(**self.training_tqdm_dict)
+
+        # collapse all metrics into one dict
+        all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()}
+
+        # track all metrics for callbacks
+        self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()}
+
+        return 0, grad_norm_dic, all_log_metrics
+
+    def training_forward(self, batch, batch_nb, opt_idx, hiddens):
+        """
+        Handle forward for each training case (distributed, single gpu, etc...)
+        :param batch:
+        :param batch_nb:
+        :return:
+        """
+        # ---------------
+        # FORWARD
+        # ---------------
+        # enable not needing to add opt_idx to training_step
+        args = [batch, batch_nb]
+        if len(self.optimizers) > 1:
+            args.append(opt_idx)
+
+        # pass hiddens if using tbptt
+        if self.truncated_bptt_steps is not None:
+            args.append(hiddens)
+
+        # distributed forward
+        if self.use_ddp or self.use_ddp2 or self.use_dp:
+            output = self.model(*args)
+
+        # single GPU forward
+        elif self.single_gpu:
+            gpu_id = 0
+            if type(self.data_parallel_device_ids) is list:
+                gpu_id = self.data_parallel_device_ids[0]
+            batch = self.transfer_batch_to_gpu(batch, gpu_id)
+            args[0] = batch
+            output = self.model.training_step(*args)
+
+        # CPU forward
+        else:
+            output = self.model.training_step(*args)
+
+        # allow any mode to define training_end
+        if self.is_overriden('training_end'):
+            model_ref = self.get_model()
+            output = model_ref.training_end(output)
+
+        # format and reduce outputs accordingly
+        output = self.process_output(output, train=True)
+
+        return output

pytorch_lightning/trainer/trainer.py

@@ -0,0 +1,478 @@
+"""
+The trainer handles all the logic for running a val loop, training loop, distributing, etc.. .
+"""
+
+import os
+import warnings
+import logging
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import tqdm
+from torch.optim.optimizer import Optimizer
+
+from pytorch_lightning.trainer.amp_mixin import TrainerAMPMixin
+from pytorch_lightning.trainer.callback_config_mixin import TrainerCallbackConfigMixin
+from pytorch_lightning.trainer.data_loading_mixin import TrainerDataLoadingMixin
+from pytorch_lightning.trainer.ddp_mixin import TrainerDDPMixin
+from pytorch_lightning.trainer.dp_mixin import TrainerDPMixin
+from pytorch_lightning.trainer.dp_mixin import (
+    parse_gpu_ids,
+    determine_root_gpu_device
+)
+from pytorch_lightning.trainer.evaluation_loop_mixin import TrainerEvaluationLoopMixin
+from pytorch_lightning.trainer.logging_mixin import TrainerLoggingMixin
+from pytorch_lightning.trainer.model_hooks_mixin import TrainerModelHooksMixin
+from pytorch_lightning.trainer.train_loop_mixin import TrainerTrainLoopMixin
+from pytorch_lightning.trainer.trainer_io import TrainerIOMixin
+from pytorch_lightning.trainer.training_tricks_mixin import TrainerTrainingTricksMixin
+from pytorch_lightning.utilities.debugging import MisconfigurationException
+
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+
+
+class Trainer(TrainerIOMixin,
+              TrainerDDPMixin,
+              TrainerDPMixin,
+              TrainerDataLoadingMixin,
+              TrainerAMPMixin,
+              TrainerEvaluationLoopMixin,
+              TrainerTrainLoopMixin,
+              TrainerLoggingMixin,
+              TrainerTrainingTricksMixin,
+              TrainerCallbackConfigMixin,
+              TrainerModelHooksMixin):
+
+    def __init__(self,
+                 logger=True,
+                 checkpoint_callback=True,
+                 early_stop_callback=True,
+                 default_save_path=None,
+                 gradient_clip_val=0,
+                 gradient_clip=None,  # backward compatible
+                 process_position=0,
+                 nb_gpu_nodes=1,
+                 gpus=None,
+                 log_gpu_memory=None,
+                 show_progress_bar=True,
+                 overfit_pct=0.0,
+                 track_grad_norm=-1,
+                 check_val_every_n_epoch=1,
+                 fast_dev_run=False,
+                 accumulate_grad_batches=1,
+                 max_nb_epochs=1000,
+                 min_nb_epochs=1,
+                 train_percent_check=1.0,
+                 val_percent_check=1.0,
+                 test_percent_check=1.0,
+                 val_check_interval=1.0,
+                 log_save_interval=100,
+                 row_log_interval=10,
+                 add_row_log_interval=None,  # backward compatible
+                 distributed_backend=None,
+                 use_amp=False,
+                 print_nan_grads=False,
+                 weights_summary='full',
+                 weights_save_path=None,
+                 amp_level='O1',
+                 nb_sanity_val_steps=5,
+                 truncated_bptt_steps=None):
+        """
+
+        :param logger: Logger for experiment tracking
+        :param checkpoint_callback: Callback for checkpointing
+        :param early_stop_callback: Callback for early stopping
+        :param default_save_path: Default path for logs+weights if no logger/ckpt_callback passed
+        :param gradient_clip_val: int. 0 means don't clip.
+        :param gradient_clip: int. 0 means don't clip. Deprecated.
+        :param process_position: shown in the tqdm bar
+        :param nb_gpu_nodes: number of GPU nodes
+        :param gpus: int. (ie: 2 gpus) OR list to specify which GPUs [0, 1] OR '0,1'
+            OR '-1' / -1 to use all available gpus
+        :param log_gpu_memory: str. None, 'min_max', 'all'
+        :param show_progress_bar: Bool. If true shows tqdm bar
+        :param overfit_pct: float. uses this much of all datasets
+        :param track_grad_norm: int. -1 no tracking. Otherwise tracks that norm
+        :param check_val_every_n_epoch: int. check val every n train epochs
+        :param fast_dev_run: Bool. runs full iteration over everything to find bugs
+        :param accumulate_grad_batches: int. Accumulates grads every k batches
+        :param max_nb_epochs: int.
+        :param min_nb_epochs: int.
+        :param train_percent_check: int. How much of train set to check
+        :param val_percent_check: int. How much of val set to check
+        :param test_percent_check: int. How much of test set to check
+        :param val_check_interval: float/int. If float, % of tng epoch. If int, check every n batch
+        :param log_save_interval: int. Writes logs to disk this often
+        :param row_log_interval: int. How often to add logging rows
+        :param add_row_log_interval: int. How often to add logging rows. Deprecated.
+        :param distributed_backend: str. Options: 'dp', 'ddp', 'ddp2'.
+        :param use_amp: Bool. If true uses apex for 16bit precision
+        :param print_nan_grads: Bool. Prints nan gradients
+        :param weights_summary: str. Options: 'full', 'top', None to not print.
+        :param weights_save_path: Bool. Where to save weights if on cluster
+        :param amp_level: str. Check nvidia docs for level
+        :param nb_sanity_val_steps: int. How many val steps before a full train loop.
+        :param truncated_bptt_steps: int. Enables multiple backward passes for each batch.
+        """
+        # Transfer params
+        self.nb_gpu_nodes = nb_gpu_nodes
+        self.log_gpu_memory = log_gpu_memory
+        if not (gradient_clip is None):
+            # Backward compatibility
+            warnings.warn("gradient_clip has renamed to gradient_clip_val since v0.5.0",
+                          DeprecationWarning)
+            gradient_clip_val = gradient_clip
+        self.gradient_clip_val = gradient_clip_val
+        self.check_val_every_n_epoch = check_val_every_n_epoch
+        self.track_grad_norm = track_grad_norm
+        self.on_gpu = gpus is not None and torch.cuda.is_available()
+        self.process_position = process_position
+        self.weights_summary = weights_summary
+        self.max_nb_epochs = max_nb_epochs
+        self.min_nb_epochs = min_nb_epochs
+        self.nb_sanity_val_steps = nb_sanity_val_steps
+        self.print_nan_grads = print_nan_grads
+        self.truncated_bptt_steps = truncated_bptt_steps
+        self.shown_warnings = set()
+
+        self.fast_dev_run = fast_dev_run
+        if self.fast_dev_run:
+            self.nb_sanity_val_steps = 1
+            self.max_nb_epochs = 1
+            m = '''
+            Running in fast_dev_run mode: will run a full train,
+            val loop using a single batch
+            '''
+            logging.info(m)
+
+        # set default save path if user didn't provide one
+        self.default_save_path = default_save_path
+        if self.default_save_path is None:
+            self.default_save_path = os.getcwd()
+
+        # training bookeeping
+        self.total_batch_nb = 0
+        self.running_loss = []
+        self.avg_loss = 0
+        self.batch_nb = 0
+        self.tqdm_metrics = {}
+        self.callback_metrics = {}
+        self.nb_val_batches = 0
+        self.nb_training_batches = 0
+        self.nb_test_batches = 0
+        self.get_train_dataloader = None
+        self.get_test_dataloaders = None
+        self.get_val_dataloaders = None
+        self.is_iterable_train_dataloader = False
+
+        # training state
+        self.model = None
+        self.testing = False
+        self.lr_schedulers = []
+        self.optimizers = None
+        self.global_step = 0
+        self.current_epoch = 0
+        self.total_batches = 0
+
+        # configure early stop callback
+        # creates a default one if none passed in
+        self.early_stop_callback = None
+        self.configure_early_stopping(early_stop_callback, logger)
+
+        # configure checkpoint callback
+        self.checkpoint_callback = checkpoint_callback
+        self.weights_save_path = weights_save_path
+
+        # accumulated grads
+        self.configure_accumulated_gradients(accumulate_grad_batches)
+
+        # allow int, string and gpu list
+        self.data_parallel_device_ids = parse_gpu_ids(gpus)
+        self.root_gpu = determine_root_gpu_device(self.data_parallel_device_ids)
+
+        # distributed backend choice
+        self.use_ddp = False
+        self.use_ddp2 = False
+        self.use_dp = False
+        self.single_gpu = False
+        self.distributed_backend = distributed_backend
+        self.set_distributed_mode(distributed_backend, nb_gpu_nodes)
+
+        # init flags for SLURM+ddp to work
+        self.proc_rank = 0
+        self.world_size = 1
+        self.node_rank = 0
+        self.configure_slurm_ddp(nb_gpu_nodes)
+
+        # nvidia setup
+        self.set_nvidia_flags(self.is_slurm_managing_tasks, self.data_parallel_device_ids)
+
+        # can't init progress bar here because starting a new process
+        # means the progress_bar won't survive pickling
+        self.show_progress_bar = show_progress_bar
+
+        # logging
+        self.log_save_interval = log_save_interval
+        self.val_check_interval = val_check_interval
+        if not (add_row_log_interval is None):
+            # backward compatibility
+            warnings.warn("gradient_clip has renamed to gradient_clip_val since v0.5.0",
+                          DeprecationWarning)
+            row_log_interval = add_row_log_interval
+        self.row_log_interval = row_log_interval
+
+        # how much of the data to use
+        self.determine_data_use_amount(train_percent_check, val_percent_check,
+                                       test_percent_check, overfit_pct)
+
+        # 16 bit mixed precision training using apex
+        self.amp_level = amp_level
+        self.init_amp(use_amp)
+
+        # set logging options
+        logging.basicConfig(level=logging.INFO)
+
+    @property
+    def slurm_job_id(self):
+        try:
+            job_id = os.environ['SLURM_JOB_ID']
+            job_id = int(job_id)
+        except Exception as e:
+            job_id = None
+        return job_id
+
+    def __parse_gpu_ids(self, gpus):
+        """
+        :param gpus: Int, string or list of ids
+        :return:
+        """
+        # if gpus = -1 then use all available devices
+        # otherwise, split the string using commas
+        if gpus is not None:
+            if type(gpus) is list:
+                gpus = gpus
+            elif type(gpus) is str:
+                if gpus == '-1':
+                    gpus = list(range(0, torch.cuda.device_count()))
+                else:
+                    gpus = [int(x.strip()) for x in gpus.split(',')]
+            elif type(gpus) is int:
+                gpus = gpus
+            else:
+                raise Exception('gpus has to be a string, int or list of ints')
+
+        return gpus
+
+    def __set_root_gpu(self, gpus):
+        if gpus is None:
+            return None
+
+        # set root gpu
+        root_gpu = 0
+        if type(gpus) is list:
+            root_gpu = gpus[0]
+
+        return root_gpu
+
+    @property
+    def num_gpus(self):
+        gpus = self.data_parallel_device_ids
+        if gpus is None:
+            return 0
+        else:
+            return len(gpus)
+
+    @property
+    def data_parallel(self):
+        return self.use_dp or self.use_ddp or self.use_ddp2
+
+    @property
+    def training_tqdm_dict(self):
+        """
+        Read-only for tqdm metrics
+        :return:
+        """
+        tqdm_dict = {
+            'loss': '{0:.3f}'.format(self.avg_loss),
+            'batch_nb': '{}'.format(self.batch_nb),
+        }
+
+        if self.truncated_bptt_steps is not None:
+            tqdm_dict['split_nb'] = self.split_nb
+
+        if self.logger is not None and self.logger.version is not None:
+            tqdm_dict['v_nb'] = self.logger.version
+
+        tqdm_dict.update(self.tqdm_metrics)
+
+        if self.on_gpu:
+            tqdm_dict['gpu'] = '{}'.format(torch.cuda.current_device())
+
+        return tqdm_dict
+
+    @property
+    def tng_tqdm_dic(self):
+        """
+        * Deprecated in v0.5.0. use training_tqdm_dict instead. *
+        :return:
+        """
+        warnings.warn("tng_tqdm_dict has renamed to training_tqdm_dict since v0.5.0",
+                      DeprecationWarning)
+        return self.training_tqdm_dict
+
+    # -----------------------------
+    # MODEL TRAINING
+    # -----------------------------
+    def fit(self, model):
+        # when using multi-node or DDP within a node start each module in a separate process
+        if self.use_ddp2:
+            task = int(os.environ['SLURM_LOCALID'])
+            self.ddp_train(task, model)
+
+        elif self.use_ddp:
+            if self.is_slurm_managing_tasks:
+                task = int(os.environ['SLURM_LOCALID'])
+                self.ddp_train(task, model)
+            else:
+                mp.spawn(self.ddp_train, nprocs=self.num_gpus, args=(model,))
+
+        # 1 gpu or dp option triggers training using DP module
+        # easier to avoid NCCL issues
+        elif self.use_dp:
+            self.dp_train(model)
+
+        elif self.single_gpu:
+            self.single_gpu_train(model)
+
+        # ON CPU
+        else:
+            # run through amp wrapper
+            if self.use_amp:
+                raise MisconfigurationException('amp + cpu is not supported.'
+                                                ' Please use a GPU option')
+
+            # CHOOSE OPTIMIZER
+            # allow for lr schedulers as well
+            self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+
+            self.run_pretrain_routine(model)
+
+        # return 1 when finished
+        # used for testing or when we need to know that training succeeded
+        return 1
+
+    def init_optimizers(self, optimizers):
+
+        # single optimizer
+        if isinstance(optimizers, Optimizer):
+            return [optimizers], []
+
+        # two lists
+        elif len(optimizers) == 2 and isinstance(optimizers[0], list):
+            optimizers, lr_schedulers = optimizers
+            return optimizers, lr_schedulers
+
+        # single list or tuple
+        elif isinstance(optimizers, list) or isinstance(optimizers, tuple):
+            return optimizers, []
+
+    def run_pretrain_routine(self, model):
+        """
+        Sanity check a few things before starting actual training
+        :param model:
+        :return:
+        """
+        ref_model = model
+        if self.data_parallel:
+            ref_model = model.module
+
+        # give model convenience properties
+        ref_model.trainer = self
+
+        # set local properties on the model
+        self.copy_trainer_model_properties(ref_model)
+
+        # link up experiment object
+        if self.logger is not None:
+            ref_model.logger = self.logger
+
+            # save exp to get started
+            if hasattr(ref_model, "hparams"):
+                self.logger.log_hyperparams(ref_model.hparams)
+
+            self.logger.save()
+
+        if self.use_ddp or self.use_ddp2:
+            dist.barrier()
+
+        # set up checkpoint callback
+        self.configure_checkpoint_callback()
+
+        # register auto-resubmit when on SLURM
+        self.register_slurm_signal_handlers()
+
+        # transfer data loaders from model
+        self.get_dataloaders(ref_model)
+
+        # print model summary
+        if self.proc_rank == 0 and self.weights_summary is not None:
+            if self.weights_summary in ['full', 'top']:
+                ref_model.summarize(mode=self.weights_summary)
+            else:
+                m = "weights_summary can be None, 'full' or 'top'"
+                raise MisconfigurationException(m)
+
+        # track model now.
+        # if cluster resets state, the model will update with the saved weights
+        self.model = model
+
+        # restore training and model before hpc call
+        self.restore_weights(model)
+
+        # when testing requested only run test and return
+        if self.testing:
+            self.run_evaluation(test=True)
+            return
+
+        # run tiny validation (if validation defined)
+        # to make sure program won't crash during val
+        ref_model.on_sanity_check_start()
+        if self.get_val_dataloaders() is not None and self.nb_sanity_val_steps > 0:
+            # init progress bars for validation sanity check
+            pbar = tqdm.tqdm(desc='Validation sanity check', total=self.nb_sanity_val_steps,
+                             leave=False, position=2 * self.process_position,
+                             disable=not self.show_progress_bar, dynamic_ncols=True, unit='batch')
+            self.main_progress_bar = pbar
+            # dummy validation progress bar
+            self.val_progress_bar = tqdm.tqdm(disable=True)
+
+            self.evaluate(model, self.get_val_dataloaders(), self.nb_sanity_val_steps, self.testing)
+
+            # close progress bars
+            self.main_progress_bar.close()
+            self.val_progress_bar.close()
+
+        # init progress bar
+        pbar = tqdm.tqdm(leave=True, position=2 * self.process_position,
+                         disable=not self.show_progress_bar, dynamic_ncols=True, unit='batch')
+        self.main_progress_bar = pbar
+
+        # clear cache before training
+        if self.on_gpu:
+            torch.cuda.empty_cache()
+
+        # CORE TRAINING LOOP
+        self.train()
+
+    def test(self, model=None):
+        self.testing = True
+        if model is not None:
+            self.fit(model)
+        else:
+            self.run_evaluation(test=True)

pytorch_lightning/trainer/trainer_io.py

@@ -0,0 +1,376 @@
+import os
+import re
+import signal
+import warnings
+from subprocess import call
+import logging
+
+import torch
+import torch.distributed as dist
+
+from pytorch_lightning.pt_overrides.override_data_parallel import (
+    LightningDistributedDataParallel, LightningDataParallel)
+
+
+class TrainerIOMixin(object):
+
+    def get_model(self):
+        is_dp_module = isinstance(self.model, (LightningDistributedDataParallel,
+                                               LightningDataParallel))
+        model = self.model.module if is_dp_module else self.model
+        return model
+
+    # --------------------
+    # CHECK-POINTING
+    # --------------------
+    def restore_weights(self, model):
+        """
+        To restore weights we have two cases.
+        First, attempt to restore hpc weights. If successful, don't restore
+        other weights.
+
+        Otherwise, try to restore actual weights
+        :param model:
+        :return:
+        """
+        # clear cache before restore
+        if self.on_gpu:
+            torch.cuda.empty_cache()
+
+        # if script called from hpc resubmit, load weights
+        did_restore_hpc_weights = self.restore_hpc_weights_if_needed(model)
+
+        # clear cache after restore
+        if self.on_gpu:
+            torch.cuda.empty_cache()
+
+        if not did_restore_hpc_weights:
+            # restore weights if same exp version
+            self.restore_state_if_checkpoint_exists(model)
+
+        # wait for all models to restore weights
+        if self.use_ddp or self.use_ddp2:
+            # wait for all processes to catch up
+            dist.barrier()
+
+        # clear cache after restore
+        if self.on_gpu:
+            torch.cuda.empty_cache()
+
+    def restore_state_if_checkpoint_exists(self, model):
+        did_restore = False
+
+        # do nothing if there's not dir or callback
+        no_ckpt_callback = (self.checkpoint_callback is None) or (not self.checkpoint_callback)
+        if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath):
+            return did_restore
+
+        # restore trainer state and model if there is a weight for this experiment
+        last_epoch = -1
+        last_ckpt_name = None
+
+        # find last epoch
+        checkpoints = os.listdir(self.checkpoint_callback.filepath)
+        for name in checkpoints:
+            # ignore hpc ckpts
+            if 'hpc_' in name:
+                continue
+
+            if '.ckpt' in name:
+                epoch = name.split('epoch_')[1]
+                epoch = int(re.sub('[^0-9]', '', epoch))
+
+                if epoch > last_epoch:
+                    last_epoch = epoch
+                    last_ckpt_name = name
+
+        # restore last checkpoint
+        if last_ckpt_name is not None:
+            last_ckpt_path = os.path.join(self.checkpoint_callback.filepath, last_ckpt_name)
+            self.restore(last_ckpt_path, self.on_gpu)
+            logging.info(f'model and trainer restored from checkpoint: {last_ckpt_path}')
+            did_restore = True
+
+        return did_restore
+
+    # --------------------
+    # HPC SIGNAL HANDLING
+    # --------------------
+    def register_slurm_signal_handlers(self):
+        # see if we're using slurm (not interactive)
+        on_slurm = False
+        try:
+            job_name = os.environ['SLURM_JOB_NAME']
+            if job_name != 'bash':
+                on_slurm = True
+        except Exception as e:
+            pass
+
+        if on_slurm:
+            logging.info('set slurm handle signals')
+            signal.signal(signal.SIGUSR1, self.sig_handler)
+            signal.signal(signal.SIGTERM, self.term_handler)
+
+    def sig_handler(self, signum, frame):
+        if self.proc_rank == 0:
+            # save weights
+            logging.info('handling SIGUSR1')
+            self.hpc_save(self.weights_save_path, self.logger)
+
+            # find job id
+            job_id = os.environ['SLURM_JOB_ID']
+            cmd = 'scontrol requeue {}'.format(job_id)
+
+            # requeue job
+            logging.info('\nrequeing job {job_id}...')
+            result = call(cmd, shell=True)
+
+            # print result text
+            if result == 0:
+                logging.info('requeued exp {job_id}')
+            else:
+                logging.info('requeue failed...')
+
+            # close experiment to avoid issues
+            self.logger.close()
+
+    def term_handler(self, signum, frame):
+        # save
+        logging.info("bypassing sigterm")
+
+    # --------------------
+    # MODEL SAVE CHECKPOINT
+    # --------------------
+    def save_checkpoint(self, filepath):
+        checkpoint = self.dump_checkpoint()
+
+        # do the actual save
+        try:
+            torch.save(checkpoint, filepath)
+        except AttributeError:
+            if 'hparams' in checkpoint:
+                del checkpoint['hparams']
+
+            torch.save(checkpoint, filepath)
+
+    def restore(self, checkpoint_path, on_gpu):
+
+        # if on_gpu:
+        #     checkpoint = torch.load(checkpoint_path)
+        # else:
+        # load on CPU first
+        checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
+
+        # load model state
+        model = self.get_model()
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
+        if on_gpu:
+            model.cuda(self.root_gpu)
+
+        # load training state (affects trainer only)
+        self.restore_training_state(checkpoint)
+
+    def dump_checkpoint(self):
+
+        checkpoint = {
+            'epoch': self.current_epoch,
+            'global_step': self.global_step
+        }
+
+        if self.checkpoint_callback is not None and self.checkpoint_callback is not False:
+            checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
+
+        if self.early_stop_callback is not None and self.checkpoint_callback is not False:
+            checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
+            checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
+
+        # save optimizers
+        optimizer_states = []
+        for i, optimizer in enumerate(self.optimizers):
+            optimizer_states.append(optimizer.state_dict())
+
+        checkpoint['optimizer_states'] = optimizer_states
+
+        # save lr schedulers
+        lr_schedulers = []
+        for i, scheduler in enumerate(self.lr_schedulers):
+            lr_schedulers.append(scheduler.state_dict())
+
+        checkpoint['lr_schedulers'] = lr_schedulers
+
+        # add the hparams and state_dict from the model
+        model = self.get_model()
+        checkpoint['state_dict'] = model.state_dict()
+        if hasattr(model, "hparams"):
+            checkpoint['hparams'] = vars(model.hparams)
+        else:
+            warnings.warn(
+                "Did not find hyperparameters at model.hparams. Saving checkpoint without"
+                " hyperparameters"
+            )
+
+        # give the model a chance to add a few things
+        model.on_save_checkpoint(checkpoint)
+
+        return checkpoint
+
+    # --------------------
+    # HPC IO
+    # --------------------
+    def restore_hpc_weights_if_needed(self, model):
+        """
+        If there is a set of hpc weights, use as signal to restore model
+        :param model:
+        :return:
+        """
+        did_restore = False
+
+        # look for hpc weights
+        folderpath = self.weights_save_path
+        if os.path.exists(folderpath):
+            files = os.listdir(folderpath)
+            hpc_weight_paths = [x for x in files if 'hpc_ckpt' in x]
+
+            # if hpc weights exist restore model
+            if len(hpc_weight_paths) > 0:
+                self.hpc_load(folderpath, self.on_gpu)
+                did_restore = True
+        return did_restore
+
+    def restore_training_state(self, checkpoint):
+        """
+        Restore trainer state.
+        Model will get its change to update
+        :param checkpoint:
+        :return:
+        """
+        if self.checkpoint_callback is not None and self.checkpoint_callback is not False:
+            self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
+
+        if self.early_stop_callback is not None and self.early_stop_callback is not False:
+            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']
+        for optimizer, opt_state in zip(self.optimizers, optimizer_states):
+            optimizer.load_state_dict(opt_state)
+
+            # move optimizer to GPU 1 weight at a time
+            # avoids OOM
+            if self.root_gpu is not None:
+                for state in optimizer.state.values():
+                    for k, v in state.items():
+                        if isinstance(v, torch.Tensor):
+                            state[k] = v.cuda(self.root_gpu)
+
+        # restore the lr schedulers
+        lr_schedulers = checkpoint['lr_schedulers']
+        for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers):
+            scheduler.load_state_dict(lrs_state)
+
+    # ----------------------------------
+    # PRIVATE OPS
+    # ----------------------------------
+    def hpc_save(self, folderpath, logger):
+        # make sure the checkpoint folder exists
+        os.makedirs(folderpath, exist_ok=True)
+
+        # save logger to make sure we get all the metrics
+        logger.save()
+
+        ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
+
+        if not os.path.exists(folderpath):
+            os.makedirs(folderpath, exist_ok=True)
+        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
+
+        # give model a chance to do something on hpc_save
+        model = self.get_model()
+        checkpoint = self.dump_checkpoint()
+
+        model.on_hpc_save(checkpoint)
+
+        # do the actual save
+        # TODO: fix for anything with multiprocess DP, DDP, DDP2
+        try:
+            torch.save(checkpoint, filepath)
+        except AttributeError:
+            if 'hparams' in checkpoint:
+                del checkpoint['hparams']
+
+            torch.save(checkpoint, filepath)
+
+        return filepath
+
+    def hpc_load(self, folderpath, on_gpu):
+        filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
+
+        # load on CPU first
+        checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
+
+        # load model state
+        model = self.get_model()
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
+
+        if self.root_gpu is not None:
+            model.cuda(self.root_gpu)
+
+        # load training state (affects trainer only)
+        self.restore_training_state(checkpoint)
+
+        # call model hook
+        model.on_hpc_load(checkpoint)
+
+        logging.info(f'restored hpc model from: {filepath}')
+
+    def max_ckpt_in_folder(self, path, name_key='ckpt_'):
+        files = os.listdir(path)
+        files = [x for x in files if name_key in x]
+        if len(files) == 0:
+            return 0
+
+        ckpt_vs = []
+        for name in files:
+            name = name.split(name_key)[-1]
+            name = re.sub('[^0-9]', '', name)
+            ckpt_vs.append(int(name))
+
+        return max(ckpt_vs)
+
+
+def load_hparams_from_tags_csv(tags_csv):
+    from argparse import Namespace
+    import pandas as pd
+
+    tags_df = pd.read_csv(tags_csv)
+    dic = tags_df.to_dict(orient='records')
+
+    ns_dict = {row['key']: convert(row['value']) for row in dic}
+
+    ns = Namespace(**ns_dict)
+    return ns
+
+
+def convert(val):
+    constructors = [int, float, str]
+
+    if type(val) is str:
+        if val.lower() == 'true':
+            return True
+        if val.lower() == 'false':
+            return False
+
+    for c in constructors:
+        try:
+            return c(val)
+        except ValueError:
+            pass
+    return val

pytorch_lightning/trainer/training_tricks_mixin.py

@@ -0,0 +1,28 @@
+import torch
+import logging
+from pytorch_lightning.callbacks import GradientAccumulationScheduler
+
+
+class TrainerTrainingTricksMixin(object):
+
+    def clip_gradients(self):
+        if self.gradient_clip_val > 0:
+            model = self.get_model()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip_val)
+
+    def print_nan_gradients(self):
+        model = self.get_model()
+        for param in model.parameters():
+            if torch.isnan(param.grad.float()).any():
+                logging.info(param, param.grad)
+
+    def configure_accumulated_gradients(self, accumulate_grad_batches):
+        self.accumulate_grad_batches = None
+
+        if isinstance(accumulate_grad_batches, dict):
+            self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches)
+        elif isinstance(accumulate_grad_batches, int):
+            schedule = {1: accumulate_grad_batches}
+            self.accumulation_scheduler = GradientAccumulationScheduler(schedule)
+        else:
+            raise TypeError("Gradient accumulation supports only int and dict types")

pytorch_lightning/utilities/arg_parse.py

@@ -7,8 +7,7 @@ import os
 
 
 def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None):
-
-    # tng, test, val check intervals
+    # training, test, val check intervals
     parser.add_argument('--eval_test_set', dest='eval_test_set', action='store_true',
                         help='true = run test set also')
     parser.add_argument('--check_val_every_n_epoch', default=1, type=int,
@@ -19,7 +18,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     parser.add_argument('--max_nb_epochs', default=200, type=int, help='cap epochs')
     parser.add_argument('--min_nb_epochs', default=2, type=int, help='min epochs')
     parser.add_argument('--train_percent_check', default=1.0, type=float,
-                        help='how much of tng set to check')
+                        help='how much of training set to check')
     parser.add_argument('--val_percent_check', default=1.0, type=float,
                         help='how much of val set to check')
     parser.add_argument('--test_percent_check', default=1.0, type=float,
@@ -29,7 +28,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
                         help='how much within 1 epoch to check val')
     parser.add_argument('--log_save_interval', default=100, type=int,
                         help='how many batches between log saves')
-    parser.add_argument('--add_log_row_interval', default=100, type=int,
+    parser.add_argument('--row_log_interval', default=100, type=int,
                         help='add log every k batches')
 
     # early stopping
@@ -40,7 +39,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
                         help='number of epochs until stop')
 
     # gradient handling
-    parser.add_argument('--gradient_clip', default=-1, type=int)
+    parser.add_argument('--gradient_clip_val', default=-1, type=int)
     parser.add_argument('--track_grad_norm', default=-1, type=int,
                         help='if > 0, will track this grad norm')
 
@@ -78,9 +77,9 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     # FAST training
     # use these settings to make sure network has no bugs without running a full dataset
     parser.add_argument('--fast_dev_run', dest='fast_dev_run', default=False, action='store_true',
-                        help='runs validation after 1 tng step')
+                        help='runs validation after 1 training step')
     parser.add_argument('--enable_tqdm', dest='enable_tqdm', default=False, action='store_true',
-                        help='false removes the prog bar')
+                        help='false removes the progress bar')
     parser.add_argument('--overfit', default=-1, type=float,
                         help='% of dataset to use with this option. float, or -1 for none')
 
@@ -93,7 +92,7 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     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')
+                        help='enables local training')
 
     # optimizer
     parser.add_argument('--lr_scheduler_milestones', default=None, type=str)

requirements.txt

@@ -1,7 +1,9 @@
 scikit-learn==0.20.2
-tqdm==4.32.1
+tqdm==4.35.0
 twine==1.13.0
 numpy==1.16.4
-torch>=1.1.0
+torch>=1.2.0
 torchvision>=0.3.0
-pandas
+pandas>=0.20.3
+test-tube>=0.6.9
+# future>=0.17.1  # required for buildins in setup.py

setup.cfg

@@ -11,12 +11,14 @@ markers =
     slow
     remote_data
     filterwarnings
+    gpus_param_tests
 
 [pycodestyle]
 ignore = E731,W504
 max-line-length = 120
 
 [coverage:report]
+# TODO: this looks suspicion, it should be reviewed
 exclude_lines =
     pragma: no cover
     def __repr__
@@ -39,13 +41,20 @@ exclude_lines =
     break
     pass
     os.makedirs
-
 omit =
     pytorch_lightning/callbacks/pt_callbacks.py
     tests/test_models.py
     pytorch_lightning/testing_models/lm_test_module.py
     pytorch_lightning/utilities/arg_parse.py
+    examples/templates
 
 [flake8]
-ignore = E731,W504,F401,F841
+# TODO: this should be 88 or 100 according PEP8
 max-line-length = 120
+exclude = .tox,*.egg,build,temp,examples/*
+select = E,W,F
+doctests = True
+verbose = 2
+# https://pep8.readthedocs.io/en/latest/intro.html#error-codes
+format = pylint
+ignore = E731,W504,F401,F841

setup.py

@@ -1,12 +1,37 @@
 #!/usr/bin/env python
 
+import os
+from io import open
 # Always prefer setuptools over distutils
 from setuptools import setup, find_packages
 
-# https://packaging.python.org/guides/single-sourcing-package-version/
+try:
+    import builtins
+except ImportError:
+    import __builtin__ as builtins
 
+# https://packaging.python.org/guides/single-sourcing-package-version/
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 
+PATH_ROOT = os.path.dirname(__file__)
+builtins.__LIGHTNING_SETUP__ = True
+
+import pytorch_lightning  # noqa: E402
+
+
+def load_requirements(path_dir=PATH_ROOT, comment_char='#'):
+    with open(os.path.join(path_dir, 'requirements.txt'), 'r') as file:
+        lines = [ln.strip() for ln in file.readlines()]
+    reqs = []
+    for ln in lines:
+        # filer all comments
+        if comment_char in ln:
+            ln = ln[:ln.index(comment_char)]
+        if ln:  # if requirement is not empty
+            reqs.append(ln)
+    return reqs
+
+
 # https://packaging.python.org/discussions/install-requires-vs-requirements /
 # keep the meta-data here for simplicity in reading this file... it's not obvious
 # what happens and to non-engineers they won't know to look in init ...
@@ -14,26 +39,22 @@ from setuptools import setup, find_packages
 # engineer specific practices
 setup(
     name='pytorch-lightning',
-    version='0.4.2',
-    description='The Keras for ML researchers using PyTorch',
-    author='William Falcon',
-    author_email='waf2107@columbia.edu',
-    url='https://github.com/williamFalcon/pytorch-lightning',
+    version=pytorch_lightning.__version__,
+    description=pytorch_lightning.__docs__,
+    author=pytorch_lightning.__author__,
+    author_email=pytorch_lightning.__author_email__,
+    url=pytorch_lightning.__homepage__,
     download_url='https://github.com/williamFalcon/pytorch-lightning',
-    license='Apache-2',
-    packages=find_packages(),
-    long_description=open('README.md', encoding='utf-8').read(),
-    long_description_content_type='text/markdown',
+    license=pytorch_lightning.__license__,
+    packages=find_packages(exclude=['examples']),
+    # long_description=open('README.md', encoding='utf-8').read(),
+    # long_description_content_type='text/markdown',
     include_package_data=True,
     zip_safe=False,
     keywords=['deep learning', 'pytorch', 'AI'],
     python_requires='>=3.6',
-    install_requires=[
-        'torch==1.2.0',
-        'tqdm',
-        'test-tube==0.6.8',
-        'pandas>=0.20.3',
-    ],
+    setup_requires=[],
+    install_requires=load_requirements(PATH_ROOT),
     classifiers=[
         'Environment :: Console',
         'Natural Language :: English',
@@ -51,7 +72,6 @@ setup(
         # Specify the Python versions you support here. In particular, ensure
         # that you indicate whether you support Python 2, Python 3 or both.
         'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.5',
         'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: 3.7',
     ],

tests/README.md

@@ -1,4 +1,7 @@
 # PyTorch-Lightning Tests
+Most PL tests train a full MNIST model under various trainer conditions (ddp, ddp2+amp, etc...).
+This provides testing for most combinations of important settings.
+The tests expect the model to perform to a reasonable degree of testing accuracy to pass.
 
 ## Running tests
 The automatic travis tests ONLY run CPU-based tests. Although these cover most of the use cases,
@@ -26,22 +29,6 @@ The GPU machine must have:
 2. [NVIDIA-apex](https://github.com/NVIDIA/apex#linux) installed.
 
 
-### test_models.py
-This file fits a tiny model on MNIST using these different set-ups.
-1. CPU only.
-2. Single GPU with DP.
-3. Multiple (2) GPUs using DP.
-3. Multiple (2) GPUs using DDP.
-3. Multiple (2) GPUs using DP + apex (for 16-bit precision).
-3. Multiple (2) GPUs using DDP + apex (for 16-bit precision).   
-
-For each set up it also tests:
-1. model saving.
-2. model loading.
-3. predicting with a loaded model.
-4. simulated save from HPC signal.
-5. simulated load from HPC signal.
-
 ## Running Coverage   
 Make sure to run coverage on a GPU machine with at least 2 GPUs and NVIDIA apex installed. 
 

tests/debug.py

@@ -1,17 +1,15 @@
-from pytorch_lightning import Trainer
-from examples import LightningTemplateModel
-from argparse import Namespace
-from test_tube import Experiment
-from pytorch_lightning.callbacks import ModelCheckpoint
-import os
-import shutil
-
-import pytorch_lightning as pl
 import torch
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
 
+import pytorch_lightning as pl
+
+
+# from test_models import assert_ok_test_acc, load_model, \
+#     clear_save_dir, get_test_tube_logger, get_hparams, init_save_dir, \
+#     init_checkpoint_callback, reset_seed, set_random_master_port
+
 
 class CoolModel(pl.LightningModule):
 
@@ -29,7 +27,7 @@ class CoolModel(pl.LightningModule):
     def training_step(self, batch, batch_nb):
         x, y = batch
         y_hat = self.forward(x)
-        return {'tng_loss': self.my_loss(y_hat, y)}
+        return {'training_loss': self.my_loss(y_hat, y)}
 
     def validation_step(self, batch, batch_nb):
         x, y = batch
@@ -44,7 +42,7 @@ class CoolModel(pl.LightningModule):
         return [torch.optim.Adam(self.parameters(), lr=0.02)]
 
     @pl.data_loader
-    def tng_dataloader(self):
+    def train_dataloader(self):
         return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
 
     @pl.data_loader
@@ -55,122 +53,53 @@ class CoolModel(pl.LightningModule):
     def test_dataloader(self):
         return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
 
-
-def get_model():
-    # set up model with these hyperparams
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    hparams = Namespace(**{'drop_prob': 0.2,
-                           'batch_size': 32,
-                           'in_features': 28 * 28,
-                           'learning_rate': 0.001 * 8,
-                           'optimizer_name': 'adam',
-                           'data_root': os.path.join(root_dir, 'mnist'),
-                           'out_features': 10,
-                           'hidden_dim': 1000})
-    model = LightningTemplateModel(hparams)
-
-    return model, hparams
-
-
-def get_exp(debug=True):
-    # set up exp object without actually saving logs
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
-    return exp
-
-
-def init_save_dir():
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    save_dir = os.path.join(root_dir, 'save_dir')
-
-    if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
-
-    os.makedirs(save_dir, exist_ok=True)
-
-    return save_dir
-
-
-def clear_save_dir():
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    save_dir = os.path.join(root_dir, 'save_dir')
-    if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
-
-
-def load_model(exp, save_dir):
-
-    # load trained model
-    tags_path = exp.get_data_path(exp.name, exp.version)
-    tags_path = os.path.join(tags_path, 'meta_tags.csv')
-
-    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
-    weights_dir = os.path.join(save_dir, checkpoints[0])
-
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
-                                                             tags_csv=tags_path, on_gpu=True)
-
-    assert trained_model is not None, 'loading model failed'
-
-    return trained_model
-
-
-def run_prediction(dataloader, trained_model):
-    # run prediction on 1 batch
-    for batch in dataloader:
-        break
-
-    x, y = batch
-    x = x.view(x.size(0), -1)
-
-    y_hat = trained_model(x)
-
-    # acc
-    labels_hat = torch.argmax(y_hat, dim=1)
-    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-    val_acc = torch.tensor(val_acc)
-    val_acc = val_acc.item()
-
-    print(val_acc)
-
-    assert val_acc > 0.70, 'this model is expected to get > 0.7 in test set (it got %f)' % val_acc
-
-
-def main():
-
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    exp = get_exp(False)
-    exp.save()
-
-    # exp file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
-
-    trainer = Trainer(
-        experiment=exp,
-        checkpoint_callback=checkpoint,
-        progress_bar=True,
-        max_nb_epochs=1,
-        gpus=[0, 1],
-        distributed_backend='dp',
-    )
-
-    model = CoolModel()
-
-    result = trainer.fit(model)
-
-    # correct result and ok accuracy
-    assert result == 1, 'amp + ddp model failed to complete'
-
-    # test model loading
-    pretrained_model = load_model(exp, save_dir)
-
-    # test model preds
-    run_prediction(model.test_dataloader, pretrained_model)
-
-    clear_save_dir()
-
-
-if __name__ == '__main__':
-    main()
+#
+# def main():
+#     reset_seed()
+#     set_random_master_port()
+#
+#     hparams = get_hparams()
+#     model = LightningTestModel(hparams)
+#
+#     save_dir = init_save_dir()
+#
+#     # exp file to get meta
+#     logger = get_test_tube_logger(False)
+#
+#     print(logger.debug)
+#
+#     # exp file to get weights
+#     checkpoint = init_checkpoint_callback(logger)
+#
+#     trainer_options = dict(
+#         show_progress_bar=False,
+#         max_nb_epochs=1,
+#         train_percent_check=0.4,
+#         val_percent_check=0.2,
+#         checkpoint_callback=checkpoint,
+#         logger=logger,
+#         gpus=[0, 1],
+#         distributed_backend='ddp'
+#     )
+#
+#     # fit model
+#     trainer = Trainer(**trainer_options)
+#     result = trainer.fit(model)
+#
+#     exp = logger.experiment
+#     print(os.listdir(exp.get_data_path(exp.name, exp.version)))
+#
+#     # correct result and ok accuracy
+#     assert result == 1, 'training failed to complete'
+#     pretrained_model = load_model(logger.experiment, save_dir,
+#                                   module_class=LightningTestModel)
+#
+#     # run test set
+#     new_trainer = Trainer(**trainer_options)
+#     new_trainer.test(pretrained_model)
+#
+#     # test we have good test accuracy
+#     clear_save_dir()
+#
+# if __name__ == '__main__':
+#     main()

tests/requirements.txt

@@ -5,4 +5,6 @@ pytest>=3.0.5
 pytest-cov
 flake8
 check-manifest
-test_tube
+test_tube
+mlflow
+comet_ml

tests/test_a_restore_models.py

@@ -0,0 +1,409 @@
+import os
+import logging
+
+import pytest
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.testing import LightningTestModel
+from . import testing_utils
+
+
+def test_running_test_pretrained_model_ddp():
+    """Verify test() on pretrained model"""
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    testing_utils.reset_seed()
+    testing_utils.set_random_master_port()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # exp file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    # exp file to get weights
+    checkpoint = testing_utils.init_checkpoint_callback(logger)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        logger=logger,
+        gpus=[0, 1],
+        distributed_backend='ddp'
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    exp = logger.experiment
+    logging.info(os.listdir(exp.get_data_path(exp.name, exp.version)))
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = testing_utils.load_model(logger.experiment,
+                                                trainer.checkpoint_callback.filepath,
+                                                module_class=LightningTestModel)
+
+    # run test set
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    for dataloader in model.test_dataloader():
+        testing_utils.run_prediction(dataloader, pretrained_model)
+
+    testing_utils.clear_save_dir()
+
+
+def test_running_test_pretrained_model():
+    testing_utils.reset_seed()
+
+    """Verify test() on pretrained model"""
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    # logger file to get weights
+    checkpoint = testing_utils.init_checkpoint_callback(logger)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        logger=logger
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = testing_utils.load_model(
+        logger.experiment, trainer.checkpoint_callback.filepath, module_class=LightningTestModel
+    )
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    testing_utils.assert_ok_test_acc(new_trainer)
+    testing_utils.clear_save_dir()
+
+
+def test_load_model_from_checkpoint():
+    testing_utils.reset_seed()
+
+    """Verify test() on pretrained model"""
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=True,
+        logger=False,
+        default_save_path=save_dir
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = LightningTestModel.load_from_checkpoint(
+        os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_1.ckpt")
+    )
+
+    # test that hparams loaded correctly
+    for k, v in vars(hparams).items():
+        assert getattr(pretrained_model.hparams, k) == v
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    testing_utils.assert_ok_test_acc(new_trainer)
+    testing_utils.clear_save_dir()
+
+
+def test_running_test_pretrained_model_dp():
+    testing_utils.reset_seed()
+
+    """Verify test() on pretrained model"""
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    # logger file to get weights
+    checkpoint = testing_utils.init_checkpoint_callback(logger)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        logger=logger,
+        gpus=[0, 1],
+        distributed_backend='dp'
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'training failed to complete'
+    pretrained_model = testing_utils.load_model(logger.experiment,
+                                                trainer.checkpoint_callback.filepath,
+                                                module_class=LightningTestModel)
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    testing_utils.assert_ok_test_acc(new_trainer)
+    testing_utils.clear_save_dir()
+
+
+def test_dp_resume():
+    """
+    Make sure DP continues training correctly
+    :return:
+    """
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=2,
+        gpus=2,
+        distributed_backend='dp',
+    )
+
+    save_dir = testing_utils.init_save_dir()
+
+    # get logger
+    logger = testing_utils.get_test_tube_logger(debug=False)
+
+    # exp file to get weights
+    # logger file to get weights
+    checkpoint = testing_utils.init_checkpoint_callback(logger)
+
+    # add these to the trainer options
+    trainer_options['logger'] = logger
+    trainer_options['checkpoint_callback'] = checkpoint
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
+    result = trainer.fit(model)
+
+    # track epoch before saving
+    real_global_epoch = trainer.current_epoch
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + dp model failed to complete'
+
+    # ---------------------------
+    # HPC LOAD/SAVE
+    # ---------------------------
+    # save
+    trainer.hpc_save(save_dir, logger)
+
+    # init new trainer
+    new_logger = testing_utils.get_test_tube_logger(version=logger.version)
+    trainer_options['logger'] = new_logger
+    trainer_options['checkpoint_callback'] = ModelCheckpoint(save_dir)
+    trainer_options['train_percent_check'] = 0.2
+    trainer_options['val_percent_check'] = 0.2
+    trainer_options['max_nb_epochs'] = 1
+    new_trainer = Trainer(**trainer_options)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_good_acc():
+        assert new_trainer.current_epoch == real_global_epoch and new_trainer.current_epoch > 0
+
+        # if model and state loaded correctly, predictions will be good even though we
+        # haven't trained with the new loaded model
+        dp_model = new_trainer.model
+        dp_model.eval()
+
+        dataloader = trainer.get_train_dataloader()
+        testing_utils.run_prediction(dataloader, dp_model, dp=True)
+
+    # new model
+    model = LightningTestModel(hparams)
+    model.on_sanity_check_start = assert_good_acc
+
+    # fit new model which should load hpc weights
+    new_trainer.fit(model)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    testing_utils.clear_save_dir()
+
+
+def test_cpu_restore_training():
+    """
+    Verify continue training session on CPU
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    test_logger_version = 10
+    logger = testing_utils.get_test_tube_logger(False, version=test_logger_version)
+
+    trainer_options = dict(
+        max_nb_epochs=2,
+        val_check_interval=0.50,
+        val_percent_check=0.2,
+        train_percent_check=0.2,
+        logger=logger,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    real_global_epoch = trainer.current_epoch
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # wipe-out trainer and model
+    # retrain with not much data... this simulates picking training back up after slurm
+    # we want to see if the weights come back correctly
+    new_logger = testing_utils.get_test_tube_logger(False, version=test_logger_version)
+    trainer_options = dict(
+        max_nb_epochs=2,
+        val_check_interval=0.50,
+        val_percent_check=0.2,
+        train_percent_check=0.2,
+        logger=new_logger,
+        checkpoint_callback=ModelCheckpoint(save_dir),
+    )
+    trainer = Trainer(**trainer_options)
+    model = LightningTestModel(hparams)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_good_acc():
+        assert trainer.current_epoch == real_global_epoch
+        assert trainer.current_epoch >= 0
+
+        # if model and state loaded correctly, predictions will be good even though we
+        # haven't trained with the new loaded model
+        trainer.model.eval()
+        for dataloader in trainer.get_val_dataloaders():
+            testing_utils.run_prediction(dataloader, trainer.model)
+
+    model.on_sanity_check_start = assert_good_acc
+
+    # by calling fit again, we trigger training, loading weights from the cluster
+    # and our hook to predict using current model before any more weight updates
+    trainer.fit(model)
+
+    testing_utils.clear_save_dir()
+
+
+def test_model_saving_loading():
+    """
+    Tests use case where trainer saves the model, and user loads it from tags independently
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        logger=logger,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # make a prediction
+    for dataloader in model.test_dataloader():
+        for batch in dataloader:
+            break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    # generate preds before saving model
+    model.eval()
+    pred_before_saving = model(x)
+
+    # save model
+    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
+    trainer.save_checkpoint(new_weights_path)
+
+    # load new model
+    tags_path = logger.experiment.get_data_path(logger.experiment.name, logger.experiment.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+    model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path,
+                                                   tags_csv=tags_path)
+    model_2.eval()
+
+    # make prediction
+    # assert that both predictions are the same
+    new_pred = model_2(x)
+    assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    testing_utils.clear_save_dir()
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/test_cpu_models.py

@@ -0,0 +1,399 @@
+import warnings
+
+import pytest
+import torch
+
+from pytorch_lightning import Trainer, data_loader
+from pytorch_lightning.callbacks import (
+    EarlyStopping,
+)
+from pytorch_lightning.testing import (
+    LightningTestModel,
+    LightningTestModelBase,
+    LightningTestMixin,
+)
+from . import testing_utils
+
+
+def test_early_stopping_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    stopping = EarlyStopping(monitor='val_loss')
+    trainer_options = dict(
+        early_stop_callback=stopping,
+        gradient_clip_val=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        show_progress_bar=True,
+        logger=testing_utils.get_test_tube_logger(),
+        train_percent_check=0.1,
+        val_percent_check=0.1
+    )
+
+    model, hparams = testing_utils.get_model()
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
+
+def test_lbfgs_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        print_nan_grads=True,
+        show_progress_bar=False,
+        weights_summary='top',
+        train_percent_check=1.0,
+        val_percent_check=0.2
+    )
+
+    model, hparams = testing_utils.get_model(use_test_model=True, lbfgs=True)
+    testing_utils.run_model_test_no_loggers(trainer_options,
+                                            model, hparams, on_gpu=False, min_acc=0.30)
+
+    testing_utils.clear_save_dir()
+
+
+def test_default_logger_callbacks_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        gradient_clip_val=1.0,
+        overfit_pct=0.20,
+        print_nan_grads=True,
+        show_progress_bar=False,
+        train_percent_check=0.01,
+        val_percent_check=0.01
+    )
+
+    model, hparams = testing_utils.get_model()
+    testing_utils.run_model_test_no_loggers(trainer_options, model, hparams, on_gpu=False)
+
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
+    testing_utils.clear_save_dir()
+
+
+def test_running_test_after_fitting():
+    """Verify test() on fitted model"""
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    # logger file to get weights
+    checkpoint = testing_utils.init_checkpoint_callback(logger)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        test_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        logger=logger
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1, 'training failed to complete'
+
+    trainer.test()
+
+    # test we have good test accuracy
+    testing_utils.assert_ok_test_acc(trainer)
+
+    testing_utils.clear_save_dir()
+
+
+def test_running_test_without_val():
+    testing_utils.reset_seed()
+
+    """Verify test() works on a model with no val_loader"""
+
+    class CurrentTestModel(LightningTestMixin, LightningTestModelBase):
+        pass
+
+    hparams = testing_utils.get_hparams()
+    model = CurrentTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    # logger file to get weights
+    checkpoint = testing_utils.init_checkpoint_callback(logger)
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        test_percent_check=0.2,
+        checkpoint_callback=checkpoint,
+        logger=logger
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1, 'training failed to complete'
+
+    trainer.test()
+
+    # test we have good test accuracy
+    testing_utils.assert_ok_test_acc(trainer)
+
+    testing_utils.clear_save_dir()
+
+
+def test_single_gpu_batch_parse():
+    testing_utils.reset_seed()
+
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    trainer = Trainer()
+
+    # batch is just a tensor
+    batch = torch.rand(2, 3)
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch.device.index == 0 and batch.type() == 'torch.cuda.FloatTensor'
+
+    # tensor list
+    batch = [torch.rand(2, 3), torch.rand(2, 3)]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0].device.index == 0 and batch[0].type() == 'torch.cuda.FloatTensor'
+    assert batch[1].device.index == 0 and batch[1].type() == 'torch.cuda.FloatTensor'
+
+    # tensor list of lists
+    batch = [[torch.rand(2, 3), torch.rand(2, 3)]]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0][0].device.index == 0 and batch[0][0].type() == 'torch.cuda.FloatTensor'
+    assert batch[0][1].device.index == 0 and batch[0][1].type() == 'torch.cuda.FloatTensor'
+
+    # tensor dict
+    batch = [{'a': torch.rand(2, 3), 'b': torch.rand(2, 3)}]
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0]['a'].device.index == 0 and batch[0]['a'].type() == 'torch.cuda.FloatTensor'
+    assert batch[0]['b'].device.index == 0 and batch[0]['b'].type() == 'torch.cuda.FloatTensor'
+
+    # tuple of tensor list and list of tensor dict
+    batch = ([torch.rand(2, 3) for _ in range(2)],
+             [{'a': torch.rand(2, 3), 'b': torch.rand(2, 3)} for _ in range(2)])
+    batch = trainer.transfer_batch_to_gpu(batch, 0)
+    assert batch[0][0].device.index == 0 and batch[0][0].type() == 'torch.cuda.FloatTensor'
+
+    assert batch[1][0]['a'].device.index == 0
+    assert batch[1][0]['a'].type() == 'torch.cuda.FloatTensor'
+
+    assert batch[1][0]['b'].device.index == 0
+    assert batch[1][0]['b'].type() == 'torch.cuda.FloatTensor'
+
+
+def test_simple_cpu():
+    """
+    Verify continue training session on CPU
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    trainer_options = dict(
+        max_nb_epochs=1,
+        val_percent_check=0.1,
+        train_percent_check=0.1,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    testing_utils.clear_save_dir()
+
+
+def test_cpu_model():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        logger=testing_utils.get_test_tube_logger(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = testing_utils.get_model()
+
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_all_features_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    trainer_options = dict(
+        gradient_clip_val=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        show_progress_bar=False,
+        logger=testing_utils.get_test_tube_logger(),
+        accumulate_grad_batches=2,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = testing_utils.get_model()
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_tbptt_cpu_model():
+    """
+    Test truncated back propagation through time works.
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    save_dir = testing_utils.init_save_dir()
+
+    truncated_bptt_steps = 2
+    sequence_size = 30
+    batch_size = 30
+
+    x_seq = torch.rand(batch_size, sequence_size, 1)
+    y_seq_list = torch.rand(batch_size, sequence_size, 1).tolist()
+
+    class MockSeq2SeqDataset(torch.utils.data.Dataset):
+        def __getitem__(self, i):
+            return x_seq, y_seq_list
+
+        def __len__(self):
+            return 1
+
+    class BpttTestModel(LightningTestModelBase):
+        def __init__(self, hparams):
+            super().__init__(hparams)
+            self.test_hidden = None
+
+        def training_step(self, batch, batch_idx, hiddens):
+            assert hiddens == self.test_hidden, "Hidden state not persistent between tbptt steps"
+            self.test_hidden = torch.rand(1)
+
+            x_tensor, y_list = batch
+            assert x_tensor.shape[1] == truncated_bptt_steps, "tbptt split Tensor failed"
+
+            y_tensor = torch.tensor(y_list, dtype=x_tensor.dtype)
+            assert y_tensor.shape[1] == truncated_bptt_steps, "tbptt split list failed"
+
+            pred = self.forward(x_tensor.view(batch_size, truncated_bptt_steps))
+            loss_val = torch.nn.functional.mse_loss(
+                pred, y_tensor.view(batch_size, truncated_bptt_steps))
+            return {
+                'loss': loss_val,
+                'hiddens': self.test_hidden,
+            }
+
+        @data_loader
+        def train_dataloader(self):
+            return torch.utils.data.DataLoader(
+                dataset=MockSeq2SeqDataset(),
+                batch_size=batch_size,
+                shuffle=False,
+                sampler=None,
+            )
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        truncated_bptt_steps=truncated_bptt_steps,
+        val_percent_check=0,
+        weights_summary=None,
+    )
+
+    hparams = testing_utils.get_hparams()
+    hparams.batch_size = batch_size
+    hparams.in_features = truncated_bptt_steps
+    hparams.hidden_dim = truncated_bptt_steps
+    hparams.out_features = truncated_bptt_steps
+
+    model = BpttTestModel(hparams)
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1, 'training failed to complete'
+
+    testing_utils.clear_save_dir()
+
+
+def test_single_gpu_model():
+    """
+    Make sure single GPU works (DP mode)
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    if not torch.cuda.is_available():
+        warnings.warn('test_single_gpu_model cannot run.'
+                      ' Rerun on a GPU node to run this test')
+        return
+    model, hparams = testing_utils.get_model()
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus=1
+    )
+
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/test_gpu_models.py

@@ -0,0 +1,407 @@
+import os
+import pytest
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import (
+    ModelCheckpoint,
+)
+from pytorch_lightning.root_module import memory
+from pytorch_lightning.testing import (
+    LightningTestModel,
+)
+from pytorch_lightning.trainer.dp_mixin import (
+    parse_gpu_ids,
+    determine_root_gpu_device,
+)
+from pytorch_lightning.utilities.debugging import MisconfigurationException
+from . import testing_utils
+
+PRETEND_N_OF_GPUS = 16
+
+
+def test_multi_gpu_model_ddp2():
+    """
+    Make sure DDP2 works
+    :return:
+    """
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    testing_utils.reset_seed()
+    testing_utils.set_random_master_port()
+
+    model, hparams = testing_utils.get_model()
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        gpus=2,
+        weights_summary=None,
+        distributed_backend='ddp2'
+    )
+
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_ddp():
+    """
+    Make sure DDP works
+    :return:
+    """
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    testing_utils.reset_seed()
+    testing_utils.set_random_master_port()
+
+    model, hparams = testing_utils.get_model()
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        gpus=[0, 1],
+        distributed_backend='ddp'
+    )
+
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_optimizer_return_options():
+    testing_utils.reset_seed()
+
+    trainer = Trainer()
+    model, hparams = testing_utils.get_model()
+
+    # single optimizer
+    opt_a = torch.optim.Adam(model.parameters(), lr=0.002)
+    opt_b = torch.optim.SGD(model.parameters(), lr=0.002)
+    optim, lr_sched = trainer.init_optimizers(opt_a)
+    assert len(optim) == 1 and len(lr_sched) == 0
+
+    # opt tuple
+    opts = (opt_a, opt_b)
+    optim, lr_sched = trainer.init_optimizers(opts)
+    assert len(optim) == 2 and optim[0] == opts[0] and optim[1] == opts[1]
+    assert len(lr_sched) == 0
+
+    # opt list
+    opts = [opt_a, opt_b]
+    optim, lr_sched = trainer.init_optimizers(opts)
+    assert len(optim) == 2 and optim[0] == opts[0] and optim[1] == opts[1]
+    assert len(lr_sched) == 0
+
+    # opt tuple of lists
+    opts = ([opt_a], ['lr_scheduler'])
+    optim, lr_sched = trainer.init_optimizers(opts)
+    assert len(optim) == 1 and len(lr_sched) == 1
+    assert optim[0] == opts[0][0] and lr_sched[0] == 'lr_scheduler'
+
+
+def test_cpu_slurm_save_load():
+    """
+    Verify model save/load/checkpoint on CPU
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    version = logger.version
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        logger=logger,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    real_global_step = trainer.global_step
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # predict with trained model before saving
+    # make a prediction
+    for dataloader in model.test_dataloader():
+        for batch in dataloader:
+            break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    model.eval()
+    pred_before_saving = model(x)
+
+    # test HPC saving
+    # simulate snapshot on slurm
+    saved_filepath = trainer.hpc_save(save_dir, logger)
+    assert os.path.exists(saved_filepath)
+
+    # new logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False, version=version)
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        logger=logger,
+        checkpoint_callback=ModelCheckpoint(save_dir),
+    )
+    trainer = Trainer(**trainer_options)
+    model = LightningTestModel(hparams)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_pred_same():
+        assert trainer.global_step == real_global_step and trainer.global_step > 0
+
+        # predict with loaded model to make sure answers are the same
+        trainer.model.eval()
+        new_pred = trainer.model(x)
+        assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    model.on_epoch_start = assert_pred_same
+
+    # by calling fit again, we trigger training, loading weights from the cluster
+    # and our hook to predict using current model before any more weight updates
+    trainer.fit(model)
+
+    testing_utils.clear_save_dir()
+
+
+def test_multi_gpu_none_backend():
+    """
+    Make sure when using multiple GPUs the user can't use
+    distributed_backend = None
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    model, hparams = testing_utils.get_model()
+    trainer_options = dict(
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus='-1'
+    )
+
+    with pytest.raises(MisconfigurationException):
+        testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_dp():
+    """
+    Make sure DP works
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    model, hparams = testing_utils.get_model()
+    trainer_options = dict(
+        show_progress_bar=False,
+        distributed_backend='dp',
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus='-1'
+    )
+
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+    # test memory helper functions
+    memory.get_memory_profile('min_max')
+
+
+def test_ddp_sampler_error():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    testing_utils.reset_seed()
+    testing_utils.set_random_master_port()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
+
+    logger = testing_utils.get_test_tube_logger(True)
+
+    trainer = Trainer(
+        logger=logger,
+        show_progress_bar=False,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    with pytest.warns(UserWarning):
+        trainer.get_dataloaders(model)
+
+    testing_utils.clear_save_dir()
+
+
+@pytest.fixture
+def mocked_device_count(monkeypatch):
+    def device_count():
+        return PRETEND_N_OF_GPUS
+
+    monkeypatch.setattr(torch.cuda, 'device_count', device_count)
+
+
+@pytest.fixture
+def mocked_device_count_0(monkeypatch):
+    def device_count():
+        return 0
+
+    monkeypatch.setattr(torch.cuda, 'device_count', device_count)
+
+
+test_num_gpus_data = [
+    pytest.param(None, 0, None, id="None - expect 0 gpu to use."),
+    pytest.param(0, 0, None, id="Oth gpu, expect 1 gpu to use."),
+    pytest.param(1, 1, None, id="1st gpu, expect 1 gpu to use."),
+    pytest.param(-1, PRETEND_N_OF_GPUS, "ddp", id="-1 - use all gpus"),
+    pytest.param('-1', PRETEND_N_OF_GPUS, "ddp", id="'-1' - use all gpus"),
+    pytest.param(3, 3, "ddp", id="3rd gpu - 1 gpu to use (backend:ddp)")
+]
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize(["gpus", "expected_num_gpus", "distributed_backend"], test_num_gpus_data)
+def test_trainer_gpu_parse(mocked_device_count, gpus, expected_num_gpus, distributed_backend):
+    assert Trainer(gpus=gpus, distributed_backend=distributed_backend).num_gpus == expected_num_gpus
+
+
+test_num_gpus_data_0 = [
+    pytest.param(None, 0, None, id="None - expect 0 gpu to use."),
+    pytest.param(None, 0, "ddp", id="None - expect 0 gpu to use."),
+]
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize(["gpus", "expected_num_gpus", "distributed_backend"], test_num_gpus_data_0)
+def test_trainer_num_gpu_0(mocked_device_count_0, gpus, expected_num_gpus, distributed_backend):
+    assert Trainer(gpus=gpus, distributed_backend=distributed_backend).num_gpus == expected_num_gpus
+
+
+test_root_gpu_data = [
+    pytest.param(None, None, "ddp", id="None is None"),
+    pytest.param(0, None, "ddp", id="O gpus, expect gpu root device to be None."),
+    pytest.param(1, 0, "ddp", id="1 gpu, expect gpu root device to be 0."),
+    pytest.param(-1, 0, "ddp", id="-1 - use all gpus, expect gpu root device to be 0."),
+    pytest.param('-1', 0, "ddp", id="'-1' - use all gpus, expect gpu root device to be 0."),
+    pytest.param(3, 0, "ddp", id="3 gpus, expect gpu root device to be 0.(backend:ddp)")]
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize(['gpus', 'expected_root_gpu', "distributed_backend"], test_root_gpu_data)
+def test_root_gpu_property(mocked_device_count, gpus, expected_root_gpu, distributed_backend):
+    assert Trainer(gpus=gpus, distributed_backend=distributed_backend).root_gpu == expected_root_gpu
+
+
+test_root_gpu_data_for_0_devices_passing = [
+    pytest.param(None, None, None, id="None is None"),
+    pytest.param(None, None, "ddp", id="None is None"),
+    pytest.param(0, None, "ddp", id="None is None"),
+]
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize([
+    'gpus', 'expected_root_gpu', "distributed_backend"], test_root_gpu_data_for_0_devices_passing)
+def test_root_gpu_property_0_passing(
+        mocked_device_count_0, gpus, expected_root_gpu, distributed_backend):
+    assert Trainer(gpus=gpus, distributed_backend=distributed_backend).root_gpu == expected_root_gpu
+
+
+# Asking for a gpu when non are available will result in a MisconfigurationException
+test_root_gpu_data_for_0_devices_raising = [
+    pytest.param(1, None, "ddp"),
+    pytest.param(3, None, "ddp"),
+    pytest.param(3, None, "ddp"),
+    pytest.param([1, 2], None, "ddp"),
+    pytest.param([0, 1], None, "ddp"),
+    pytest.param(-1, None, "ddp"),
+    pytest.param('-1', None, "ddp")
+]
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize([
+    'gpus', 'expected_root_gpu', "distributed_backend"], test_root_gpu_data_for_0_devices_raising)
+def test_root_gpu_property_0_raising(
+        mocked_device_count_0, gpus, expected_root_gpu, distributed_backend):
+    with pytest.raises(MisconfigurationException):
+        Trainer(gpus=gpus, distributed_backend=distributed_backend).root_gpu
+
+
+test_determine_root_gpu_device_data = [
+    pytest.param(None, None, id="No gpus, expect gpu root device to be None"),
+    pytest.param([0], 0, id="Oth gpu, expect gpu root device to be 0."),
+    pytest.param([1], 1, id="1st gpu, expect gpu root device to be 1."),
+    pytest.param([3], 3, id="3rd gpu, expect gpu root device to be 3."),
+    pytest.param([1, 2], 1, id="[1, 2] gpus, expect gpu root device to be 1."),
+]
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize(['gpus', 'expected_root_gpu'], test_determine_root_gpu_device_data)
+def test_determine_root_gpu_device(gpus, expected_root_gpu):
+    assert determine_root_gpu_device(gpus) == expected_root_gpu
+
+
+test_parse_gpu_ids_data = [
+    pytest.param(None, None),
+    pytest.param(0, None),
+    pytest.param(1, [0]),
+    pytest.param(-1, list(range(PRETEND_N_OF_GPUS)), id="-1 - use all gpus"),
+    pytest.param('-1', list(range(PRETEND_N_OF_GPUS)), id="'-1' - use all gpus"),
+    pytest.param(3, [0, 1, 2])]
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize(['gpus', 'expected_gpu_ids'], test_parse_gpu_ids_data)
+def test_parse_gpu_ids(mocked_device_count, gpus, expected_gpu_ids):
+    assert parse_gpu_ids(gpus) == expected_gpu_ids
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize("gpus", [[1, 2, 19], -1, '-1'])
+def test_parse_gpu_fail_on_non_existant_id(mocked_device_count_0, gpus):
+    with pytest.raises(MisconfigurationException):
+        parse_gpu_ids(gpus)
+
+
+@pytest.mark.gpus_param_tests
+def test_parse_gpu_fail_on_non_existant_id_2(mocked_device_count):
+    with pytest.raises(MisconfigurationException):
+        parse_gpu_ids([1, 2, 19])
+
+
+@pytest.mark.gpus_param_tests
+@pytest.mark.parametrize("gpus", [-1, '-1'])
+def test_parse_gpu_returns_None_when_no_devices_are_available(mocked_device_count_0, gpus):
+    with pytest.raises(MisconfigurationException):
+        parse_gpu_ids(gpus)
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/test_models.py

@@ -1,823 +0,0 @@
-import os
-import shutil
-import warnings
-from argparse import Namespace
-
-import pytest
-import numpy as np
-import torch
-from test_tube import Experiment, SlurmCluster
-
-# sys.path += [os.path.abspath('..'), os.path.abspath('../..')]
-from pytorch_lightning import Trainer
-from pytorch_lightning.testing.lm_test_module import LightningTestModel
-from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
-from pytorch_lightning.utilities.debugging import MisconfigurationException
-from pytorch_lightning.root_module import memory
-from pytorch_lightning.models.trainer import reduce_distributed_output
-from pytorch_lightning.root_module import model_saving
-from examples import LightningTemplateModel
-
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
-
-
-# ------------------------------------------------------------------------
-# TESTS
-# ------------------------------------------------------------------------
-def test_simple_cpu():
-    """
-    Verify continue training session on CPU
-    :return:
-    """
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    trainer_options = dict(
-        max_nb_epochs=1,
-        val_percent_check=0.1,
-        train_percent_check=0.1,
-    )
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    result = trainer.fit(model)
-
-    # traning complete
-    assert result == 1, 'amp + ddp model failed to complete'
-
-    clear_save_dir()
-
-
-def test_amp_single_gpu():
-    """
-    Make sure DDP + AMP work
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a node with 2+ GPUs to run this test')
-        return
-
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    trainer_options = dict(
-        progress_bar=True,
-        max_nb_epochs=1,
-        gpus=[0],
-        distributed_backend='dp',
-        use_amp=True
-    )
-
-    run_gpu_model_test(trainer_options, model, hparams)
-
-
-def test_cpu_restore_training():
-    """
-    Verify continue training session on CPU
-    :return:
-    """
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    test_exp_version = 10
-    exp = get_exp(False, version=test_exp_version)
-    exp.argparse(hparams)
-    exp.save()
-
-    trainer_options = dict(
-        max_nb_epochs=2,
-        val_check_interval=0.50,
-        val_percent_check=0.2,
-        train_percent_check=0.2,
-        experiment=exp,
-        checkpoint_callback=ModelCheckpoint(save_dir)
-    )
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    result = trainer.fit(model)
-    real_global_epoch = trainer.current_epoch
-
-    # traning complete
-    assert result == 1, 'amp + ddp model failed to complete'
-
-    # wipe-out trainer and model
-    # retrain with not much data... this simulates picking training back up after slurm
-    # we want to see if the weights come back correctly
-    new_exp = get_exp(False, version=test_exp_version)
-    trainer_options = dict(
-        max_nb_epochs=2,
-        val_check_interval=0.50,
-        val_percent_check=0.2,
-        train_percent_check=0.2,
-        experiment=new_exp,
-        checkpoint_callback=ModelCheckpoint(save_dir),
-    )
-    trainer = Trainer(**trainer_options)
-    model = LightningTestModel(hparams)
-
-    # set the epoch start hook so we can predict before the model does the full training
-    def assert_good_acc():
-        assert trainer.current_epoch == real_global_epoch and trainer.current_epoch > 0
-
-        # if model and state loaded correctly, predictions will be good even though we
-        # haven't trained with the new loaded model
-        trainer.model.eval()
-        run_prediction(trainer.val_dataloader, trainer.model)
-
-    model.on_sanity_check_start = assert_good_acc
-
-    # by calling fit again, we trigger training, loading weights from the cluster
-    # and our hook to predict using current model before any more weight updates
-    trainer.fit(model)
-
-    clear_save_dir()
-
-
-def test_amp_gpu_ddp():
-    """
-    Make sure DDP + AMP work
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      'Rerun on a node with 2+ GPUs to run this test')
-        return
-
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    trainer_options = dict(
-        progress_bar=True,
-        max_nb_epochs=1,
-        gpus=[0, 1],
-        distributed_backend='ddp',
-        use_amp=True
-    )
-
-    run_gpu_model_test(trainer_options, model, hparams)
-
-
-def test_cpu_slurm_save_load():
-    """
-    Verify model save/load/checkpoint on CPU
-    :return:
-    """
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    exp = get_exp(False)
-    exp.argparse(hparams)
-    exp.save()
-
-    cluster_a = SlurmCluster()
-    trainer_options = dict(
-        max_nb_epochs=1,
-        cluster=cluster_a,
-        experiment=exp,
-        checkpoint_callback=ModelCheckpoint(save_dir)
-    )
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    result = trainer.fit(model)
-    real_global_step = trainer.global_step
-
-    # traning complete
-    assert result == 1, 'amp + ddp model failed to complete'
-
-    # predict with trained model before saving
-    # make a prediction
-    for batch in model.test_dataloader:
-        break
-
-    x, y = batch
-    x = x.view(x.size(0), -1)
-
-    model.eval()
-    pred_before_saving = model(x)
-
-    # test registering a save function
-    trainer.enable_auto_hpc_walltime_manager()
-
-    # test HPC saving
-    # simulate snapshot on slurm
-    saved_filepath = trainer.hpc_save(save_dir, exp)
-    assert os.path.exists(saved_filepath)
-
-    # wipe-out trainer and model
-    # retrain with not much data... this simulates picking training back up after slurm
-    # we want to see if the weights come back correctly
-    continue_tng_hparams = get_hparams(continue_training=True,
-                                       hpc_exp_number=cluster_a.hpc_exp_number)
-    trainer_options = dict(
-        max_nb_epochs=1,
-        cluster=SlurmCluster(continue_tng_hparams),
-        experiment=exp,
-        checkpoint_callback=ModelCheckpoint(save_dir),
-    )
-    trainer = Trainer(**trainer_options)
-    model = LightningTestModel(hparams)
-
-    # set the epoch start hook so we can predict before the model does the full training
-    def assert_pred_same():
-        assert trainer.global_step == real_global_step and trainer.global_step > 0
-
-        # predict with loaded model to make sure answers are the same
-        trainer.model.eval()
-        new_pred = trainer.model(x)
-        assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
-
-    model.on_epoch_start = assert_pred_same
-
-    # by calling fit again, we trigger training, loading weights from the cluster
-    # and our hook to predict using current model before any more weight updates
-    trainer.fit(model)
-
-    clear_save_dir()
-
-
-def test_loading_meta_tags():
-    hparams = get_hparams()
-
-    # save tags
-    exp = get_exp(False)
-    exp.tag({'some_str': 'a_str', 'an_int': 1, 'a_float': 2.0})
-    exp.argparse(hparams)
-    exp.save()
-
-    # load tags
-    tags_path = exp.get_data_path(exp.name, exp.version) + '/meta_tags.csv'
-    tags = model_saving.load_hparams_from_tags_csv(tags_path)
-
-    assert tags.batch_size == 32 and tags.hidden_dim == 1000
-
-    clear_save_dir()
-
-
-def test_dp_output_reduce():
-
-    # test identity when we have a single gpu
-    out = torch.rand(3, 1)
-    assert reduce_distributed_output(out, nb_gpus=1) is out
-
-    # average when we have multiples
-    assert reduce_distributed_output(out, nb_gpus=2) == out.mean()
-
-    # when we have a dict of vals
-    out = {
-        'a': out,
-        'b': {
-            'c': out
-        }
-    }
-    reduced = reduce_distributed_output(out, nb_gpus=3)
-    assert reduced['a'] == out['a']
-    assert reduced['b']['c'] == out['b']['c']
-
-
-def test_model_saving_loading():
-    """
-    Tests use case where trainer saves the model, and user loads it from tags independently
-    :return:
-    """
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    exp = get_exp(False)
-    exp.argparse(hparams)
-    exp.save()
-
-    trainer_options = dict(
-        max_nb_epochs=1,
-        cluster=SlurmCluster(),
-        experiment=exp,
-        checkpoint_callback=ModelCheckpoint(save_dir)
-    )
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    result = trainer.fit(model)
-
-    # traning complete
-    assert result == 1, 'amp + ddp model failed to complete'
-
-    # make a prediction
-    for batch in model.test_dataloader:
-        break
-
-    x, y = batch
-    x = x.view(x.size(0), -1)
-
-    # generate preds before saving model
-    model.eval()
-    pred_before_saving = model(x)
-
-    # save model
-    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
-    trainer.save_checkpoint(new_weights_path)
-
-    # load new model
-    tags_path = exp.get_data_path(exp.name, exp.version)
-    tags_path = os.path.join(tags_path, 'meta_tags.csv')
-    model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path,
-                                                   tags_csv=tags_path, on_gpu=False)
-    model_2.eval()
-
-    # make prediction
-    # assert that both predictions are the same
-    new_pred = model_2(x)
-    assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
-
-    clear_save_dir()
-
-
-def test_model_freeze_unfreeze():
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    model.freeze()
-    model.unfreeze()
-
-
-def test_amp_gpu_ddp_slurm_managed():
-    """
-    Make sure DDP + AMP work
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
-        return
-
-    # simulate setting slurm flags
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-    os.environ['SLURM_LOCALID'] = str(0)
-
-    hparams = get_hparams()
-    model = LightningTestModel(hparams)
-
-    trainer_options = dict(
-        progress_bar=True,
-        max_nb_epochs=1,
-        gpus=[0],
-        distributed_backend='ddp',
-        use_amp=True
-    )
-
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    exp = get_exp(False)
-    exp.argparse(hparams)
-    exp.save()
-
-    # exp file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
-
-    # add these to the trainer options
-    trainer_options['checkpoint_callback'] = checkpoint
-    trainer_options['experiment'] = exp
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    trainer.is_slurm_managing_tasks = True
-    result = trainer.fit(model)
-
-    # correct result and ok accuracy
-    assert result == 1, 'amp + ddp model failed to complete'
-
-    # test root model address
-    assert trainer.resolve_root_node_address('abc') == 'abc'
-    assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
-    assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
-    assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
-
-    # test model loading with a map_location
-    map_location = 'cuda:1'
-    pretrained_model = load_model(exp, save_dir, True, map_location)
-
-    # test model preds
-    run_prediction(model.test_dataloader, pretrained_model)
-
-    if trainer.use_ddp:
-        # on hpc this would work fine... but need to hack it for the purpose of the test
-        trainer.model = pretrained_model
-        trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
-
-    # test HPC loading / saving
-    trainer.hpc_save(save_dir, exp)
-    trainer.hpc_load(save_dir, on_gpu=True)
-
-    # test freeze on gpu
-    model.freeze()
-    model.unfreeze()
-
-    clear_save_dir()
-
-
-def test_early_stopping_cpu_model():
-    """
-    Test each of the trainer options
-    :return:
-    """
-
-    stopping = EarlyStopping()
-    trainer_options = dict(
-        early_stop_callback=stopping,
-        gradient_clip=1.0,
-        overfit_pct=0.20,
-        track_grad_norm=2,
-        print_nan_grads=True,
-        progress_bar=False,
-        experiment=get_exp(),
-        train_percent_check=0.1,
-        val_percent_check=0.1
-    )
-
-    model, hparams = get_model()
-    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
-
-    # test freeze on cpu
-    model.freeze()
-    model.unfreeze()
-
-
-def test_cpu_model_with_amp():
-    """
-    Make sure model trains on CPU
-    :return:
-    """
-
-    trainer_options = dict(
-        progress_bar=False,
-        experiment=get_exp(),
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.4,
-        use_amp=True
-    )
-
-    model, hparams = get_model()
-
-    with pytest.raises((MisconfigurationException, ModuleNotFoundError)):
-        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
-
-
-def test_cpu_model():
-    """
-    Make sure model trains on CPU
-    :return:
-    """
-
-    trainer_options = dict(
-        progress_bar=False,
-        experiment=get_exp(),
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.4
-    )
-
-    model, hparams = get_model()
-
-    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
-
-
-def test_all_features_cpu_model():
-    """
-    Test each of the trainer options
-    :return:
-    """
-
-    trainer_options = dict(
-        gradient_clip=1.0,
-        overfit_pct=0.20,
-        track_grad_norm=2,
-        print_nan_grads=True,
-        progress_bar=False,
-        experiment=get_exp(),
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.4
-    )
-
-    model, hparams = get_model()
-    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
-
-
-def test_single_gpu_model():
-    """
-    Make sure single GPU works (DP mode)
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_single_gpu_model cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    model, hparams = get_model()
-
-    trainer_options = dict(
-        progress_bar=False,
-        max_nb_epochs=1,
-        train_percent_check=0.1,
-        val_percent_check=0.1,
-        gpus=[0]
-    )
-
-    run_gpu_model_test(trainer_options, model, hparams)
-
-
-def test_multi_gpu_model_dp():
-    """
-    Make sure DP works
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_multi_gpu_model_dp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_multi_gpu_model_dp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
-        return
-    model, hparams = get_model()
-    trainer_options = dict(
-        progress_bar=False,
-        max_nb_epochs=1,
-        train_percent_check=0.1,
-        val_percent_check=0.1,
-        gpus='-1'
-    )
-
-    run_gpu_model_test(trainer_options, model, hparams)
-
-    # test memory helper functions
-    memory.get_gpu_memory_map()
-
-
-def test_amp_gpu_dp():
-    """
-    Make sure DP + AMP work
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_dp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_dp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
-        return
-    model, hparams = get_model()
-    trainer_options = dict(
-        max_nb_epochs=1,
-        gpus='0, 1',  # test init with gpu string
-        distributed_backend='dp',
-        use_amp=True
-    )
-    with pytest.raises(MisconfigurationException):
-        run_gpu_model_test(trainer_options, model, hparams)
-
-
-def test_multi_gpu_model_ddp():
-    """
-    Make sure DDP works
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_multi_gpu_model_ddp cannot run.'
-                      ' Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_multi_gpu_model_ddp cannot run.'
-                      ' Rerun on a node with 2+ GPUs to run this test')
-        return
-
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-    model, hparams = get_model()
-    trainer_options = dict(
-        progress_bar=False,
-        max_nb_epochs=1,
-        train_percent_check=0.4,
-        val_percent_check=0.2,
-        gpus=[0, 1],
-        distributed_backend='ddp'
-    )
-
-    run_gpu_model_test(trainer_options, model, hparams)
-
-
-def test_ddp_sampler_error():
-    """
-    Make sure DDP + AMP work
-    :return:
-    """
-    if not torch.cuda.is_available():
-        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
-        return
-    if not torch.cuda.device_count() > 1:
-        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
-        return
-
-    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
-
-    hparams = get_hparams()
-    model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
-
-    exp = get_exp(True)
-    exp.save()
-
-    trainer = Trainer(
-        experiment=exp,
-        progress_bar=False,
-        max_nb_epochs=1,
-        gpus=[0, 1],
-        distributed_backend='ddp',
-        use_amp=True
-    )
-
-    with pytest.raises(MisconfigurationException):
-        trainer.get_dataloaders(model)
-
-    clear_save_dir()
-
-
-# ------------------------------------------------------------------------
-# UTILS
-# ------------------------------------------------------------------------
-def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
-    save_dir = init_save_dir()
-
-    # exp file to get meta
-    exp = get_exp(False)
-    exp.argparse(hparams)
-    exp.save()
-
-    # exp file to get weights
-    checkpoint = ModelCheckpoint(save_dir)
-
-    # add these to the trainer options
-    trainer_options['checkpoint_callback'] = checkpoint
-    trainer_options['experiment'] = exp
-
-    # fit model
-    trainer = Trainer(**trainer_options)
-    result = trainer.fit(model)
-
-    # correct result and ok accuracy
-    assert result == 1, 'amp + ddp model failed to complete'
-
-    # test model loading
-    pretrained_model = load_model(exp, save_dir, on_gpu)
-
-    # test model preds
-    run_prediction(model.test_dataloader, pretrained_model)
-
-    if trainer.use_ddp:
-        # on hpc this would work fine... but need to hack it for the purpose of the test
-        trainer.model = pretrained_model
-        trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
-
-    # test HPC loading / saving
-    trainer.hpc_save(save_dir, exp)
-    trainer.hpc_load(save_dir, on_gpu=on_gpu)
-
-    clear_save_dir()
-
-
-def get_hparams(continue_training=False, hpc_exp_number=0):
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-
-    args = {
-        'drop_prob': 0.2,
-        'batch_size': 32,
-        'in_features': 28 * 28,
-        'learning_rate': 0.001 * 8,
-        'optimizer_name': 'adam',
-        'data_root': os.path.join(root_dir, 'mnist'),
-        'out_features': 10,
-        'hidden_dim': 1000}
-
-    if continue_training:
-        args['test_tube_do_checkpoint_load'] = True
-        args['hpc_exp_number'] = hpc_exp_number
-
-    hparams = Namespace(**args)
-    return hparams
-
-
-def get_model(use_test_model=False):
-    # set up model with these hyperparams
-    hparams = get_hparams()
-
-    if use_test_model:
-        model = LightningTestModel(hparams)
-    else:
-        model = LightningTemplateModel(hparams)
-
-    return model, hparams
-
-
-def get_exp(debug=True, version=None):
-    # set up exp object without actually saving logs
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir', version=version)
-    return exp
-
-
-def init_save_dir():
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    save_dir = os.path.join(root_dir, 'save_dir')
-
-    if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
-
-    os.makedirs(save_dir, exist_ok=True)
-
-    return save_dir
-
-
-def clear_save_dir():
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    save_dir = os.path.join(root_dir, 'save_dir')
-    if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
-
-
-def load_model(exp, save_dir, on_gpu, map_location=None):
-
-    # load trained model
-    tags_path = exp.get_data_path(exp.name, exp.version)
-    tags_path = os.path.join(tags_path, 'meta_tags.csv')
-
-    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
-    weights_dir = os.path.join(save_dir, checkpoints[0])
-
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
-                                                             tags_csv=tags_path,
-                                                             on_gpu=on_gpu,
-                                                             map_location=map_location)
-
-    assert trained_model is not None, 'loading model failed'
-
-    return trained_model
-
-
-def run_prediction(dataloader, trained_model):
-    # run prediction on 1 batch
-    for batch in dataloader:
-        break
-
-    x, y = batch
-    x = x.view(x.size(0), -1)
-
-    y_hat = trained_model(x)
-
-    # acc
-    labels_hat = torch.argmax(y_hat, dim=1)
-    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-    val_acc = torch.tensor(val_acc)
-    val_acc = val_acc.item()
-
-    print(val_acc)
-
-    assert val_acc > 0.50, 'this model is expected to get > 0.50 in test set (it got %f)' % val_acc
-
-
-def assert_ok_acc(trainer):
-    # this model should get 0.80+ acc
-    acc = trainer.tng_tqdm_dic['val_acc']
-    assert acc > 0.50, 'model failed to get expected 0.50 validation accuracy. Got: %f' % acc
-
-
-if __name__ == '__main__':
-    pytest.main([__file__])

tests/test_trainer.py

@@ -0,0 +1,324 @@
+import os
+import pytest
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import (
+    ModelCheckpoint,
+)
+from pytorch_lightning.testing import (
+    LightningTestModel,
+    LightningTestModelBase,
+    LightningValidationStepMixin,
+    LightningValidationMultipleDataloadersMixin,
+    LightningTestMixin,
+    LightningTestMultipleDataloadersMixin,
+)
+from pytorch_lightning.trainer import trainer_io
+from pytorch_lightning.trainer.logging_mixin import TrainerLoggingMixin
+from . import testing_utils
+
+
+def test_no_val_module():
+    """
+    Tests use case where trainer saves the model, and user loads it from tags independently
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+
+    class CurrentTestModel(LightningTestModelBase):
+        pass
+
+    model = CurrentTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        logger=logger,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # training complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # save model
+    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
+    trainer.save_checkpoint(new_weights_path)
+
+    # load new model
+    tags_path = logger.experiment.get_data_path(logger.experiment.name, logger.experiment.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+    model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path,
+                                                   tags_csv=tags_path)
+    model_2.eval()
+
+    # make prediction
+    testing_utils.clear_save_dir()
+
+
+def test_no_val_end_module():
+    """
+    Tests use case where trainer saves the model, and user loads it from tags independently
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    class CurrentTestModel(LightningValidationStepMixin, LightningTestModelBase):
+        pass
+
+    hparams = testing_utils.get_hparams()
+    model = CurrentTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    # logger file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        logger=logger,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # save model
+    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
+    trainer.save_checkpoint(new_weights_path)
+
+    # load new model
+    tags_path = logger.experiment.get_data_path(logger.experiment.name, logger.experiment.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+    model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path,
+                                                   tags_csv=tags_path)
+    model_2.eval()
+
+    # make prediction
+    testing_utils.clear_save_dir()
+
+
+def test_gradient_accumulation_scheduling():
+    testing_utils.reset_seed()
+
+    """
+    Test grad accumulation by the freq of optimizer updates
+    """
+    # test incorrect configs
+    with pytest.raises(IndexError):
+        assert Trainer(accumulate_grad_batches={0: 3, 1: 4, 4: 6})
+        assert Trainer(accumulate_grad_batches={-2: 3})
+
+    with pytest.raises(TypeError):
+        assert Trainer(accumulate_grad_batches={})
+        assert Trainer(accumulate_grad_batches=[[2, 3], [4, 6]])
+        assert Trainer(accumulate_grad_batches={1: 2, 3.: 4})
+        assert Trainer(accumulate_grad_batches={1: 2.5, 3: 5})
+
+    # test optimizer call freq matches scheduler
+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, second_order_closure=None):
+        # only test the first 12 batches in epoch
+        if batch_nb < 12:
+            if epoch_nb == 0:
+                # reset counter when starting epoch
+                if batch_nb == 0:
+                    self.prev_called_batch_nb = 0
+
+                    # use this opportunity to test once
+                    assert self.trainer.accumulate_grad_batches == 1
+
+                assert batch_nb == self.prev_called_batch_nb
+                self.prev_called_batch_nb += 1
+
+            elif 1 <= epoch_nb <= 2:
+                # reset counter when starting epoch
+                if batch_nb == 1:
+                    self.prev_called_batch_nb = 1
+
+                    # use this opportunity to test once
+                    assert self.trainer.accumulate_grad_batches == 2
+
+                assert batch_nb == self.prev_called_batch_nb
+                self.prev_called_batch_nb += 2
+
+            else:
+                if batch_nb == 3:
+                    self.prev_called_batch_nb = 3
+
+                    # use this opportunity to test once
+                    assert self.trainer.accumulate_grad_batches == 4
+
+                assert batch_nb == self.prev_called_batch_nb
+                self.prev_called_batch_nb += 3
+
+        optimizer.step()
+
+        # clear gradients
+        optimizer.zero_grad()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+    schedule = {1: 2, 3: 4}
+
+    trainer = Trainer(accumulate_grad_batches=schedule,
+                      train_percent_check=0.1,
+                      val_percent_check=0.1,
+                      max_nb_epochs=4)
+
+    # for the test
+    trainer.optimizer_step = optimizer_step
+    model.prev_called_batch_nb = 0
+
+    trainer.fit(model)
+
+
+def test_loading_meta_tags():
+    testing_utils.reset_seed()
+
+    from argparse import Namespace
+    hparams = testing_utils.get_hparams()
+
+    # save tags
+    logger = testing_utils.get_test_tube_logger(False)
+    logger.log_hyperparams(Namespace(some_str='a_str', an_int=1, a_float=2.0))
+    logger.log_hyperparams(hparams)
+    logger.save()
+
+    # load tags
+    tags_path = logger.experiment.get_data_path(
+        logger.experiment.name, logger.experiment.version
+    ) + '/meta_tags.csv'
+    tags = trainer_io.load_hparams_from_tags_csv(tags_path)
+
+    assert tags.batch_size == 32 and tags.hidden_dim == 1000
+
+    testing_utils.clear_save_dir()
+
+
+def test_dp_output_reduce():
+    mixin = TrainerLoggingMixin()
+    testing_utils.reset_seed()
+
+    # test identity when we have a single gpu
+    out = torch.rand(3, 1)
+    assert mixin.reduce_distributed_output(out, nb_gpus=1) is out
+
+    # average when we have multiples
+    assert mixin.reduce_distributed_output(out, nb_gpus=2) == out.mean()
+
+    # when we have a dict of vals
+    out = {
+        'a': out,
+        'b': {
+            'c': out
+        }
+    }
+    reduced = mixin.reduce_distributed_output(out, nb_gpus=3)
+    assert reduced['a'] == out['a']
+    assert reduced['b']['c'] == out['b']['c']
+
+
+def test_model_freeze_unfreeze():
+    testing_utils.reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    model.freeze()
+    model.unfreeze()
+
+
+def test_multiple_val_dataloader():
+    """
+    Verify multiple val_dataloader
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    class CurrentTestModel(
+        LightningValidationMultipleDataloadersMixin,
+        LightningTestModelBase
+    ):
+        pass
+
+    hparams = testing_utils.get_hparams()
+    model = CurrentTestModel(hparams)
+
+    # logger file to get meta
+    trainer_options = dict(
+        max_nb_epochs=1,
+        val_percent_check=0.1,
+        train_percent_check=1.0,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # verify training completed
+    assert result == 1
+
+    # verify there are 2 val loaders
+    assert len(trainer.get_val_dataloaders()) == 2, \
+        'Multiple val_dataloaders not initiated properly'
+
+    # make sure predictions are good for each val set
+    for dataloader in trainer.get_val_dataloaders():
+        testing_utils.run_prediction(dataloader, trainer.model)
+
+
+def test_multiple_test_dataloader():
+    """
+    Verify multiple test_dataloader
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    class CurrentTestModel(
+        LightningTestMultipleDataloadersMixin,
+        LightningTestModelBase
+    ):
+        pass
+
+    hparams = testing_utils.get_hparams()
+    model = CurrentTestModel(hparams)
+
+    # logger file to get meta
+    trainer_options = dict(
+        max_nb_epochs=1,
+        val_percent_check=0.1,
+        train_percent_check=0.1,
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # verify there are 2 val loaders
+    assert len(trainer.get_test_dataloaders()) == 2, \
+        'Multiple test_dataloaders not initiated properly'
+
+    # make sure predictions are good for each test set
+    for dataloader in trainer.get_test_dataloaders():
+        testing_utils.run_prediction(dataloader, trainer.model)
+
+    # run the test method
+    trainer.test()
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/test_y_logging.py

@@ -0,0 +1,192 @@
+import os
+import pickle
+
+import numpy as np
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.testing import LightningTestModel
+from pytorch_lightning.logging import LightningLoggerBase, rank_zero_only
+from . import testing_utils
+
+RANDOM_FILE_PATHS = list(np.random.randint(12000, 19000, 1000))
+ROOT_SEED = 1234
+torch.manual_seed(ROOT_SEED)
+np.random.seed(ROOT_SEED)
+RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
+
+
+def test_testtube_logger():
+    """
+    verify that basic functionality of test tube logger works
+    """
+    reset_seed()
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    logger = testing_utils.get_test_tube_logger(False)
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        train_percent_check=0.01,
+        logger=logger
+    )
+
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    assert result == 1, "Training failed"
+
+    testing_utils.clear_save_dir()
+
+
+def test_testtube_pickle():
+    """
+    Verify that pickling a trainer containing a test tube logger works
+    """
+    reset_seed()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = testing_utils.init_save_dir()
+
+    logger = testing_utils.get_test_tube_logger(False)
+    logger.log_hyperparams(hparams)
+    logger.save()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        train_percent_check=0.01,
+        logger=logger
+    )
+
+    trainer = Trainer(**trainer_options)
+    pkl_bytes = pickle.dumps(trainer)
+    trainer2 = pickle.loads(pkl_bytes)
+    trainer2.logger.log_metrics({"acc": 1.0})
+
+    testing_utils.clear_save_dir()
+
+
+def test_mlflow_logger():
+    """
+    verify that basic functionality of mlflow logger works
+    """
+    reset_seed()
+
+    try:
+        from pytorch_lightning.logging import MLFlowLogger
+    except ModuleNotFoundError:
+        return
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    mlflow_dir = os.path.join(root_dir, "mlruns")
+
+    logger = MLFlowLogger("test", f"file://{mlflow_dir}")
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        train_percent_check=0.01,
+        logger=logger
+    )
+
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    print('result finished')
+    assert result == 1, "Training failed"
+
+    testing_utils.clear_save_dir()
+
+
+def test_mlflow_pickle():
+    """
+    verify that pickling trainer with mlflow logger works
+    """
+    reset_seed()
+
+    try:
+        from pytorch_lightning.logging import MLFlowLogger
+    except ModuleNotFoundError:
+        return
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    mlflow_dir = os.path.join(root_dir, "mlruns")
+
+    logger = MLFlowLogger("test", f"file://{mlflow_dir}")
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        logger=logger
+    )
+
+    trainer = Trainer(**trainer_options)
+    pkl_bytes = pickle.dumps(trainer)
+    trainer2 = pickle.loads(pkl_bytes)
+    trainer2.logger.log_metrics({"acc": 1.0})
+
+    testing_utils.clear_save_dir()
+
+
+def test_custom_logger(tmpdir):
+
+    class CustomLogger(LightningLoggerBase):
+        def __init__(self):
+            super().__init__()
+            self.hparams_logged = None
+            self.metrics_logged = None
+            self.finalized = False
+
+        @rank_zero_only
+        def log_hyperparams(self, params):
+            self.hparams_logged = params
+
+        @rank_zero_only
+        def log_metrics(self, metrics, step_num):
+            self.metrics_logged = metrics
+
+        @rank_zero_only
+        def finalize(self, status):
+            self.finalized_status = status
+
+        @property
+        def name(self):
+            return "name"
+
+        @property
+        def version(self):
+            return "1"
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    logger = CustomLogger()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        train_percent_check=0.01,
+        logger=logger,
+        default_save_path=tmpdir
+    )
+
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    assert result == 1, "Training failed"
+    assert logger.hparams_logged == hparams
+    assert logger.metrics_logged != {}
+    assert logger.finalized_status == "success"
+
+
+def reset_seed():
+    SEED = RANDOM_SEEDS.pop()
+    torch.manual_seed(SEED)
+    np.random.seed(SEED)

tests/test_z_amp.py

@@ -0,0 +1,208 @@
+import os
+import warnings
+
+import pytest
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.testing import (
+    LightningTestModel,
+)
+from pytorch_lightning.utilities.debugging import MisconfigurationException
+from . import testing_utils
+
+
+def test_amp_single_gpu():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=1,
+        gpus=1,
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_no_amp_single_gpu():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=1,
+        gpus=1,
+        distributed_backend='dp',
+        use_amp=True
+    )
+
+    with pytest.raises((MisconfigurationException, ModuleNotFoundError)):
+        testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_amp_gpu_ddp():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    testing_utils.reset_seed()
+    testing_utils.set_random_master_port()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=1,
+        gpus=2,
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_amp_gpu_ddp_slurm_managed():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    testing_utils.reset_seed()
+
+    # simulate setting slurm flags
+    testing_utils.set_random_master_port()
+    os.environ['SLURM_LOCALID'] = str(0)
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        show_progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    save_dir = testing_utils.init_save_dir()
+
+    # exp file to get meta
+    logger = testing_utils.get_test_tube_logger(False)
+
+    # exp file to get weights
+    checkpoint = testing_utils.init_checkpoint_callback(logger)
+
+    # add these to the trainer options
+    trainer_options['checkpoint_callback'] = checkpoint
+    trainer_options['logger'] = logger
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test root model address
+    assert trainer.resolve_root_node_address('abc') == 'abc'
+    assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
+    assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
+    assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
+
+    # test model loading with a map_location
+    pretrained_model = testing_utils.load_model(logger.experiment,
+                                                trainer.checkpoint_callback.filepath)
+
+    # test model preds
+    for dataloader in trainer.get_test_dataloaders():
+        testing_utils.run_prediction(dataloader, pretrained_model)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, logger)
+    trainer.hpc_load(save_dir, on_gpu=True)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    testing_utils.clear_save_dir()
+
+
+def test_cpu_model_with_amp():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    trainer_options = dict(
+        show_progress_bar=False,
+        logger=testing_utils.get_test_tube_logger(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4,
+        use_amp=True
+    )
+
+    model, hparams = testing_utils.get_model()
+
+    with pytest.raises((MisconfigurationException, ModuleNotFoundError)):
+        testing_utils.run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_amp_gpu_dp():
+    """
+    Make sure DP + AMP work
+    :return:
+    """
+    testing_utils.reset_seed()
+
+    if not testing_utils.can_run_gpu_test():
+        return
+
+    model, hparams = testing_utils.get_model()
+    trainer_options = dict(
+        max_nb_epochs=1,
+        gpus='0, 1',  # test init with gpu string
+        distributed_backend='dp',
+        use_amp=True
+    )
+    with pytest.raises(MisconfigurationException):
+        testing_utils.run_gpu_model_test(trainer_options, model, hparams)
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

tests/testing_utils.py

@@ -0,0 +1,239 @@
+import os
+import shutil
+import warnings
+from argparse import Namespace
+
+import numpy as np
+import torch
+
+from pl_examples import LightningTemplateModel
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import (
+    ModelCheckpoint,
+)
+from pytorch_lightning.logging import TestTubeLogger
+from pytorch_lightning.testing import (
+    LightningTestModel,
+)
+
+# generate a list of random seeds for each test
+RANDOM_FILE_PATHS = list(np.random.randint(12000, 19000, 1000))
+RANDOM_PORTS = list(np.random.randint(12000, 19000, 1000))
+ROOT_SEED = 1234
+torch.manual_seed(ROOT_SEED)
+np.random.seed(ROOT_SEED)
+RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
+
+
+def run_model_test_no_loggers(trainer_options, model, hparams, on_gpu=True, min_acc=0.50):
+    save_dir = init_save_dir()
+    trainer_options['default_save_path'] = save_dir
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(trainer.logger.experiment,
+                                  trainer.checkpoint_callback.filepath)
+
+    # test new model accuracy
+    for dataloader in model.test_dataloader():
+        run_prediction(dataloader, pretrained_model, min_acc=min_acc)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
+
+    clear_save_dir()
+
+
+def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
+    save_dir = init_save_dir()
+
+    # logger file to get meta
+    logger = get_test_tube_logger(False)
+
+    # logger file to get weights
+    checkpoint = init_checkpoint_callback(logger)
+
+    # add these to the trainer options
+    trainer_options['checkpoint_callback'] = checkpoint
+    trainer_options['logger'] = logger
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(logger.experiment, trainer.checkpoint_callback.filepath)
+
+    # test new model accuracy
+    [run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
+
+    if trainer.use_ddp or trainer.use_ddp2:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, logger)
+    trainer.hpc_load(save_dir, on_gpu=on_gpu)
+
+    clear_save_dir()
+
+
+def get_hparams(continue_training=False, hpc_exp_number=0):
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+
+    args = {
+        'drop_prob': 0.2,
+        'batch_size': 32,
+        'in_features': 28 * 28,
+        'learning_rate': 0.001 * 8,
+        'optimizer_name': 'adam',
+        'data_root': os.path.join(root_dir, 'mnist'),
+        'out_features': 10,
+        'hidden_dim': 1000}
+
+    if continue_training:
+        args['test_tube_do_checkpoint_load'] = True
+        args['hpc_exp_number'] = hpc_exp_number
+
+    hparams = Namespace(**args)
+    return hparams
+
+
+def get_model(use_test_model=False, lbfgs=False):
+    # set up model with these hyperparams
+    hparams = get_hparams()
+    if lbfgs:
+        setattr(hparams, 'optimizer_name', 'lbfgs')
+        setattr(hparams, 'learning_rate', 0.002)
+
+    if use_test_model:
+        model = LightningTestModel(hparams)
+    else:
+        model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_test_tube_logger(debug=True, version=None):
+    # set up logger object without actually saving logs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+    logger = TestTubeLogger(save_dir, name='lightning_logs', debug=False, version=version)
+    return logger
+
+
+def init_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'tests', 'save_dir')
+
+    if os.path.exists(save_dir):
+        n = RANDOM_FILE_PATHS.pop()
+        shutil.move(save_dir, save_dir + f'_{n}')
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+def clear_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+    if os.path.exists(save_dir):
+        n = RANDOM_FILE_PATHS.pop()
+        shutil.move(save_dir, save_dir + f'_{n}')
+
+
+def load_model(exp, root_weights_dir, module_class=LightningTemplateModel):
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(root_weights_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(root_weights_dir, checkpoints[0])
+
+    trained_model = module_class.load_from_metrics(weights_path=weights_dir,
+                                                   tags_csv=tags_path)
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model, dp=False, min_acc=0.50):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    if dp:
+        output = trained_model(batch, 0)
+        acc = output['val_acc']
+        acc = torch.mean(acc).item()
+
+    else:
+        y_hat = trained_model(x)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        acc = torch.tensor(acc)
+        acc = acc.item()
+
+    assert acc > min_acc, f'this model is expected to get > {min_acc} in test set (it got {acc})'
+
+
+def assert_ok_val_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.training_tqdm_dict['val_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def assert_ok_test_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.training_tqdm_dict['test_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+def can_run_gpu_test():
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_ddp cannot run.'
+                      ' Rerun on a GPU node to run this test')
+        return False
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_ddp cannot run.'
+                      ' Rerun on a node with 2+ GPUs to run this test')
+        return False
+    return True
+
+
+def reset_seed():
+    SEED = RANDOM_SEEDS.pop()
+    torch.manual_seed(SEED)
+    np.random.seed(SEED)
+
+
+def set_random_master_port():
+    port = RANDOM_PORTS.pop()
+    os.environ['MASTER_PORT'] = str(port)
+
+
+def init_checkpoint_callback(logger):
+    exp = logger.experiment
+    exp_path = exp.get_data_path(exp.name, exp.version)
+    ckpt_dir = os.path.join(exp_path, 'checkpoints')
+    checkpoint = ModelCheckpoint(ckpt_dir)
+    return checkpoint

tox.ini

@@ -12,36 +12,40 @@
 #  and also to help confirm pull requests to this project.
 
 [tox]
-envlist = py{35,36,37}
+envlist = py{35,36,37,38}
 
-[pytest]
-log_cli = 0
-log_cli_level = CRITICAL
-log_cli_format = %(message)s
-log_file = pytest.log
-log_file_level = DEBUG
-log_file_format = %(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)
-log_file_date_format=%Y-%m-%d %H:%M:%S
+# DROP, it is duplication of setup.cfg
+# [pytest]
+# log_cli = 0
+# log_cli_level = CRITICAL
+# log_cli_format = %(message)s
+# log_file = pytest.log
+# log_file_level = DEBUG
+# log_file_format = %(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)
+# log_file_date_format=%Y-%m-%d %H:%M:%S
 
 [testenv]
 basepython =
     py35: python3.5
     py36: python3.6
     py37: python3.7
+    py38: python3.8
 deps =
     -r requirements.txt
     -r ./tests/requirements.txt
 commands =
+    pip list
     check-manifest --ignore tox.ini
-    python setup.py check -m -s
-    coverage run --source pytorch_lightning -m py.test pytorch_lightning tests examples -v --doctest-modules
+    python setup.py check --metadata --strict
+    coverage run --source pytorch_lightning -m py.test pytorch_lightning tests pl_examples -v --doctest-modules
     flake8 .
 
-[flake8]
-exclude = .tox,*.egg,build,temp
-select = E,W,F
-doctests = True
-verbose = 2
+# DROP, it is duplication of setup.cfg
+# [flake8]
+# exclude = .tox,*.egg,build,temp,examples/*
+# select = E,W,F
+# doctests = True
+# verbose = 2
 # https://pep8.readthedocs.io/en/latest/intro.html#error-codes
-format = pylint
-max-line-length = 100
+# format = pylint
+# max-line-length = 100