release v0.5.1.2

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

* cleaning up demos

.codecov.yml

@@ -0,0 +1,46 @@
+# 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
+
+coverage:
+  precision: 0  # 2 = xx.xx%, 0 = xx%
+  round: nearest # how coverage is rounded: down/up/nearest
+  range: 40...100 # custom range of coverage colors from red -> yellow -> green
+  status:
+    # https://codecov.readme.io/v1.0/docs/commit-status
+    project:
+      default:
+        against: auto
+        target: 99% # specify the target coverage for each commit status
+        threshold: 20% # allow this little decrease on project
+        # https://github.com/codecov/support/wiki/Filtering-Branches
+        # branches: master
+        if_ci_failed: error
+    # https://github.com/codecov/support/wiki/Patch-Status
+    patch:
+      default:
+        against: auto
+        target: 40% # specify the target "X%" coverage to hit
+        # threshold: 50% # allow this much decrease on patch
+    changes: false
+
+parsers:
+  gcov:
+    branch_detection:
+      conditional: true
+      loop: true
+      macro: false
+      method: false
+  javascript:
+    enable_partials: false
+
+comment:
+  layout: header, diff
+  require_changes: false
+  behavior: default  # update if exists else create new
+  # branches: *
+

.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

@@ -0,0 +1,42 @@
+# 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!   
+
+## 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   
+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    
+It's useful for users to look at the code and understand very quickly what's happening. Many users won't be engineers. Thus we need to value clear, simple code over condensed ninja moves. While that's super cool, this isn't the project for that :)      
+
+#### Force User Decisions To Best Practices    
+There are 1,000 ways to do something. However, something eventually becomes standard practice that everyone does. Thus we pick one way of doing it and force everyone to do it this way. A good example is accumulated gradients. There are many ways to implement, we just pick one and force users to use that one. A bad forced decision would be to make users use a specific library to do something.    
+
+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.    
+
+#### 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    
+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:  
+1. Submit a github issue.   
+2. Fix it.  
+3. Submit a PR! 
+
+## New Features:  
+1. Submit a github issue.   
+2. We'll agree on the feature scope.     
+3. Submit a PR! (with updated docs and tests 🙃).   

.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,36 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: bug
+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.
+
+**To Reproduce**
+Steps to reproduce the behavior:
+1. Go to '...'
+2. Click on '....'
+3. Scroll down to '....'
+4. See error
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+
+**Desktop (please complete the following information):**
+ - OS: [e.g. iOS]
+ - Browser [e.g. chrome, safari]
+ - Version [e.g. 22]
+
+**Additional context**
+Add any other context about the problem here.

.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: enhancement, help wanted
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request 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/ISSUE_TEMPLATE/typos-and-doc-fixes.md

@@ -0,0 +1,17 @@
+---
+name: Typos and doc fixes
+about: Typos and doc fixes
+title: ''
+labels: typo
+assignees: ''
+
+---
+
+For typos and doc fixes, please go ahead and:
+
+1. Create an issue.
+2. Fix the typo.   
+3. Submit a PR.
+
+
+Thanks!

.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

@@ -9,6 +9,9 @@ model_weights/
 app/models/
 pip-wheel-metadata/
 test_tube_exp/
+tests/tests_tt_dir/
+tests/save_dir
+default/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -119,4 +122,4 @@ ENV/
 .mypy_cache/
 
 # data
-mnist/
+mnist/

.readthedocs.yml

@@ -0,0 +1,19 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation with MkDocs
+mkdocs:
+  configuration: mkdocs.yml
+
+# Optionally build your docs in additional formats such as PDF and ePub
+formats: all
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  version: 3.7
+  install:
+    - requirements: docs/requirements.txt

.run_local_tests.sh

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

.travis.yml

@@ -0,0 +1,70 @@
+# vim ft=yaml
+
+# After changing this file, check it on:
+#   http://yaml-online-parser.appspot.com/
+
+# See doc/travis_notes.txt for some guidelines
+
+# this file is *not* meant to cover or endorse the use of travis, but rather to
+# help confirm pull requests to this project.
+
+env:
+  global:
+    - DISPLAY=""
+
+language: python
+
+matrix:
+  include:
+    - os: linux
+      dist: xenial  # Ubuntu 16.04
+      python: 3.6
+      env: TOXENV=py36
+    - 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
+      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
+
+script:
+  # integration
+  - tox --sitepackages
+  - python setup.py install --dry-run
+
+after_success:
+  - coverage report
+  # disable auto coverage bc it isn't accurate since it misses gpu code.
+  # to get coverage, run local and push results
+  # - codecov
+
+notifications:
+  email: false

COPYING

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2019 William Falcon
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

MANIFEST.in

@@ -1,9 +1,42 @@
-graft docs
+# Manifest syntax https://docs.python.org/2/distutils/sourcedist.html
+graft wheelhouse
 
-include COPYING
-include AUTHORS
+recursive-include birl *.py
+recursive-exclude __pycache__  *.py[cod] *.orig
 
-recursive-include src/einsteinpy/tests *.py *.html
+# Include the README
+include *.md
 
-prune docs/source/examples/.ipynb_checkpoints
-global-exclude *.py[cod] __pycache__ *.so *.dylib
+# Include the license file
+include LICENSE
+
+exclude *.sh
+exclude *.toml
+exclude *.svg    
+recursive-include pytorch_lightning *.py
+
+# include examples
+recursive-include examples *.py
+recursive-include examples *.md
+recursive-include examples *.sh
+
+# exclude tests from package
+recursive-exclude tests *
+recursive-exclude site *
+exclude tests
+
+# Exclude the documentation files
+recursive-exclude docs *
+exclude docs
+
+# Include the Requirements
+include requirements.txt
+
+# Exclude build configs
+exclude *.yml
+
+prune .git
+prune .github
+prune notebook*
+prune temp*
+prune test*

README.md

@@ -1,56 +1,195 @@
-<p align="center">
-  <a href="https://williamfalcon.github.io/pytorch-lightning/">
-    <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/lightning_logo.png" width="50">
-  </a>
-</p>
-<h3 align="center">
-  Pytorch Lightning
-</h3>
-<p align="center">
-  The Keras for ML researchers using PyTorch. More control. Less boilerplate.    
-</p>
-<p align="center">
-  <a href="https://badge.fury.io/py/pytorch-lightning"><img src="https://badge.fury.io/py/pytorch-lightning.svg" alt="PyPI version" height="18"></a>
-<!--   <a href="https://travis-ci.org/williamFalcon/test-tube"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a> -->
-  <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/COPYING"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
-</p>   
+<div align="center">
 
+![Logo](./docs/source/_static/lightning_logo_small.png)
+
+# PyTorch Lightning
+
+**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)
+[![Build status](https://ci.appveyor.com/api/projects/status/rum89d7hq8l1kfye?svg=true)](https://ci.appveyor.com/project/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)
+[![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)
+<!-- 
+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)
+-->
+
+</div>
+
+Simple installation from PyPI
 ```bash
-pip install pytorch-lightning    
+pip install pytorch-lightning  
 ```
 
 ## Docs   
 **[View the docs here](https://williamfalcon.github.io/pytorch-lightning/)**
 
 ## What is it?  
-Keras and fast.ai are too abstract for researchers. Lightning abstracts the full training loop but gives you control in the critical points.   
+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?
-Because you don't want to define a training loop, validation loop, gradient clipping, checkpointing, loading, 
-gpu training, etc... every time you start a project. Let lightning handle all of that for you! Just define your 
-data and what happens in the training, testing and validation loop and lightning will do the rest.
+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.   
+
+Lightning sets up all the boilerplate state-of-the-art training for you so you can focus on the research.   
+
+---
+ 
+## 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)    
+- [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)    
+- [Bleeding edge install](https://github.com/williamFalcon/pytorch-lightning#bleeding-edge)   
+- [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 Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py).   
-2. [Define a LightningModel](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py).     
+1. [Define a LightningModule](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)     
 
-## What does lightning control for me?
-Everything!    
-Except for these 6 core functions which you define:    
+**WARNING:** This syntax is for version 0.5.0+ where abbreviations were removed.    
+```python
+import os
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
 
-```{.python}
+import pytorch_lightning as pl
+
+class CoolSystem(pl.LightningModule):
+
+    def __init__(self):
+        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 training_step(self, batch, batch_nb):
+        # REQUIRED
+        x, y = batch
+        y_hat = self.forward(x)
+        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': F.cross_entropy(y_hat, y)}
+
+    def validation_end(self, outputs):
+        # OPTIONAL
+        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+        tensorboard_logs = {'val_loss': avg_loss}
+        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}
+
+    def configure_optimizers(self):
+        # 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 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):
+        # OPTIONAL
+        return DataLoader(MNIST(os.getcwd(), train=True, 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 = CoolSystem()
+
+# most basic trainer, uses good defaults
+trainer = Trainer()    
+trainer.fit(model)   
+```     
+
+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   
+# train on cpu using only 10% of the data (for demo purposes)
+trainer = Trainer(max_nb_epochs=1, train_percent_check=0.1)
+
+# train on 4 gpus (lightning chooses GPUs for you)
+# trainer = Trainer(max_nb_epochs=1, gpus=4)  
+
+# train on 4 gpus (you choose GPUs)
+# trainer = Trainer(max_nb_epochs=1, gpus=[0, 1, 3, 7])   
+
+# train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)
+# trainer = Trainer(max_nb_epochs=1, gpus=8, nb_gpu_nodes=4)
+
+# train (1 epoch only here for demo)
+trainer.fit(model)
+
+# view tensorboard logs 
+print('View tensorboard logs by running\ntensorboard --logdir %s' % os.getcwd())
+print('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:  
+
+![Overview](./docs/source/_static/overview_flat.jpg)
+
+```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():
 ```
@@ -59,13 +198,13 @@ 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)
      
-    # even as complex as a seq-2seq + attn model
+    # even as complex as a seq-2-seq + attn model
     # (this is just a toy, non-working example to illustrate)
     start_token = '<SOS>'
     last_hidden = torch.zeros(...)
@@ -86,8 +225,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)
@@ -112,65 +251,46 @@ 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.   
 
-<p align="center">
-  <a href="https://williamfalcon.github.io/pytorch-lightning/">
-    <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/tf_loss.png" width="900px">
-  </a>
-</p>
+![tensorboard-support](./docs/source/_static/tf_loss.png)
 
 Lightning also adds a text column with all the hyperparameters for this experiment.      
 
-<p align="center">
-  <a href="https://williamfalcon.github.io/pytorch-lightning/">
-        <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/tf_tags.png" width="900px">
-  </a>
-</p>
+![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
+## Lightning automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
 
-exp = Experiment(save_dir='/some/path')
-trainer = Trainer(experiment=exp)
-...
-```   
-
-And run tensorboard from that dir   
-```bash
-tensorboard --logdir /some/path     
-```    
-
-## Lightning automatically automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
-
-###### Checkpointing    
+#### 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)
 - [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    
+#### Distributed training    
 
 - [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)
@@ -179,55 +299,119 @@ 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 arbitrary metrics
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
 - [Process position](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#process-position)
+- [Tensorboard support](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#tensorboard-support)
 - [Save a snapshot of all hyperparameters](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#save-a-snapshot-of-all-hyperparameters) 
 - [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)
-- [Anneal Learning rate](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#anneal-learning-rate)
 - [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)
-- [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/Pytorch-Lightning/LightningModule/#configure_optimizers)
+- [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/)
 - [Set how much of the validation set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-validation-set-to-check)
 - [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/)  
 
-## Demo
+## Examples   
+- [GAN](https://github.com/williamFalcon/pytorch-lightning/tree/master/examples/domain_templates/gan.py)    
+- [MNIST](https://github.com/williamFalcon/pytorch-lightning/tree/master/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/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.   
+2. Reproducibility. If all projects use the LightningModule template, it will be much much easier to understand what's going on and where to look! It will also mean every implementation follows a standard format.   
+3. Democratizing PyTorch power user features. Distributed training? 16-bit? know you need them but don't want to take the time to implement? All good... these come built into Lightning.    
+
+**How does Lightning compare with Ignite and fast.ai?**     
+[Here's a thorough comparison](https://medium.com/@_willfalcon/pytorch-lightning-vs-pytorch-ignite-vs-fast-ai-61dc7480ad8a).    
+
+**Is this another library I have to learn?**    
+Nope! We use pure Pytorch everywhere and don't add unecessary abstractions!   
+
+**Are there plans to support Python 2?**          
+Nope.     
+
+**Are there plans to support virtualenv?**    
+Nope. Please use anaconda or miniconda.    
+
+**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**
+    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
-# install lightning
-pip install pytorch-lightning
-
-# clone lightning for the demo
-git clone https://github.com/williamFalcon/pytorch-lightning.git
-cd examples/new_project_templates/
-
-# run demo (on cpu)
-python trainer_gpu_cluster_template.py
+pip install https://github.com/williamFalcon/pytorch-lightning/archive/0.4.4.zip --upgrade
 ```
-
-Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.
-```bash
-# run a grid search on two gpus
-python fully_featured_trainer.py --gpus "0;1"
-
-# run single model on multiple gpus
-python fully_featured_trainer.py --gpus "0;1" --interactive
-```
-
-

appveyor.yml

@@ -0,0 +1,64 @@
+# https://www.appveyor.com/docs/appveyor-yml/
+environment:
+
+  # SDK v7.0 MSVC Express 2008's SetEnv.cmd script will fail if the
+  # /E:ON and /V:ON options are not enabled in the batch script interpreter
+  # See: http://stackoverflow.com/a/13751649/163740
+  CMD_IN_ENV: "cmd /E:ON /V:ON /C obvci_appveyor_python_build_env.cmd"
+
+  matrix:
+    # Pre-installed Python versions, which Appveyor may upgrade to
+    # a later point release.
+    # See: http://www.appveyor.com/docs/installed-software#python
+
+
+    # - PYTHON: "C:\\Python35-x64"
+    #   PYTHON_VERSION: "3.5.x"
+    #   PYTHON_ARCH: "64"
+    #  TOXENV: "py35"
+
+    - PYTHON: "C:\\Python36-x64"
+      PYTHON_VERSION: "3.6.x"
+      PYTHON_ARCH: "64"
+      TOXENV: "py36"
+      PIP_PYVER: "36"
+
+    - PYTHON: "C:\\Python37-x64"
+      PYTHON_VERSION: "3.7.x"
+      PYTHON_ARCH: "64"
+      TOXENV: "py37"
+      PIP_PYVER: "37"
+
+build: off
+
+# https://www.appveyor.com/docs/build-cache/
+cache:
+  - C:\ProgramData\chocolatey\bin -> appveyor.yml
+  - C:\ProgramData\chocolatey\lib -> appveyor.yml
+  - '%LOCALAPPDATA%\pip\Cache -> appveyor.yml'
+
+# scripts that run after cloning repository
+install:
+  # If there is a newer build queued for the same PR, cancel this one.
+  # The AppVeyor 'rollout builds' option is supposed to serve the same
+  # purpose but it is problematic because it tends to cancel builds pushed
+  # directly to master instead of just PR builds (or the converse).
+  - SET PATH=%PYTHON%;%PYTHON%\\Scripts;%path%
+  - pip install -U --user pip
+  - 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")
+before_test:
+  - python --version
+  - pip --version
+  - pip list
+  - dir
+
+# to run your custom scripts instead of automatic tests
+test_script:
+  - tox --sitepackages --parallel auto
+
+on_success:
+  - coverage report
+  # - codecov

docs/LightningModule/RequiredTrainerInterface.md

@@ -3,35 +3,114 @@
 
 A lightning module is a strict superclass of nn.Module, it provides a standard interface for the trainer to interact with the model.
 
-The easiest thing to do is copy [this template](../../examples/new_project_templates/lightning_module_template.py) and modify accordingly. 
+The easiest thing to do is copy the [minimal example](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#minimal-example) below and modify accordingly. 
 
 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)
-- [get_save_dict](RequiredTrainerInterface.md#get_save_dict)
-- [load_model_specific](RequiredTrainerInterface.md#load_model_specific)
-
-- [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**:   
 
-- [update_tng_log_metrics](RequiredTrainerInterface.md#update_tng_log_metrics)
-- [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)
+- [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   
+```python
+import os
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
+
+import pytorch_lightning as pl
+
+class CoolModel(pl.LightningModule):
+
+    def __init__(self):
+        super(CoolModel, 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 training_step(self, batch, batch_nb):
+        # REQUIRED
+        x, y = batch
+        y_hat = self.forward(x)
+        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': 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):
+        # REQUIRED
+        return torch.optim.Adam(self.parameters(), lr=0.02)
+
+    @pl.data_loader
+    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):
+        # 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?     
+The LightningModule interface is on the right. Each method corresponds to a part of a research project. Lightning automates everything not in blue.    
+
+<p align="center">
+  <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/overview_flat.jpg">
+    <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/overview_flat.jpg" height="900px">
+  </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.
@@ -40,7 +119,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**   
@@ -50,14 +129,15 @@ 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)
@@ -65,46 +145,147 @@ def training_step(self, data_batch, batch_nb):
     
     output = {
         'loss': loss, # required
-        'prog': {'tng_loss': loss, 'batch_nb': batch_nb} # optional
+        'progress_bar': {'training_loss': loss}, # optional (MUST ALL BE TENSORS)
+        'log': {'training_loss': loss} # optional (MUST ALL BE TENSORS)
     }
     
     # 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    
+```    
+
+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)
@@ -119,31 +300,49 @@ 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)
-```
+```   
+If you didn't define a validation_step, this won't be called.       
 
-Called at the end of the validation loop with the output of each validation_step.
+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.    
 
 **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):
     """
@@ -159,68 +358,229 @@ 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
 
-Set up as many optimizers 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.  If you use 16 bit precision it will also handle that.
-
-
-##### Return
-List - List of optimizers
-
-**Example**
-
-``` {.python}
-# most cases
-def configure_optimizers(self):
-    opt = Adam(lr=0.01)
-    return [opt]
+    val_loss_mean /= i
+    val_acc_mean /= i
+    tqdm_dict = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
     
-# gan example
-def configure_optimizers(self):
-    generator_opt = Adam(lr=0.01)
-    disriminator_opt = Adam(lr=0.02)
-    return [generator_opt, disriminator_opt] 
+    # 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
 ```
 
---- 
-### get_save_dict 
+### test_step
 
 ``` {.python}
-def get_save_dict(self)
-```
-Called by lightning to checkpoint your model. Lightning saves current epoch, current batch nb, etc...
-All you have to return is what specifically about your lightning model you want to checkpoint.
+# if you have one test dataloader:
+def test_step(self, batch, batch_nb)   
 
-##### Return
-Dictionary - No required keys. Most of the time as described in this example.   
+# 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}
-def get_save_dict(self):
-    # 99% of use cases this is all you need to return
-    checkpoint = {'state_dict': self.state_dict()}
-    return checkpoint
+# CASE 1: A single test dataset
+def test_step(self, batch, batch_nb):
+    x, y = batch
+    
+    # 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
 ```
 
 --- 
-### load_model_specific 
+### on_save_checkpoint 
 
 ``` {.python}
-def load_model_specific(self, checkpoint)
+def on_save_checkpoint(self, checkpoint)
 ```
-Called by lightning to restore your model. This is your chance to restore your model using the keys you added in get_save_dict.
-Lightning will automatically restore current epoch, batch nb, etc. 
+Called by lightning to checkpoint your model. Lightning saves the training state (current epoch, global_step, etc)
+and also saves the model state_dict. If you want to save anything else, use this method to add your own
+key-value pair.
+
+##### Return
+Nothing
+
+**Example**
+
+``` {.python}
+def on_save_checkpoint(self, checkpoint):
+    # 99% of use cases you don't need to implement this method 
+    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object
+```
+
+--- 
+### on_load_checkpoint 
+
+``` {.python}
+def on_load_checkpoint(self, checkpoint)
+```
+Called by lightning to restore your model. Lighting auto-restores global step, epoch, etc...
+It also restores the model state_dict.
+If you saved something with **on_save_checkpoint** this is your chance to restore this.
 
 ##### Return
 Nothing  
@@ -228,128 +588,81 @@ Nothing
 **Example**
 
 ``` {.python}
-def load_model_specific(self, checkpoint):
-    # you defined 'state_dict' in get_save_dict()
-    self.load_state_dict(checkpoint['state_dict'])
-```
-
---- 
-### tng_dataloader 
-
-``` {.python}
-@property
-def tng_dataloader(self)
-```
-Called by lightning during training loop. Define it as a property.
-
-##### Return
-Pytorch DataLoader
-
-**Example**
-
-``` {.python}
-@property
-def tng_dataloader(self):
-    if self._tng_dataloader is None:
-        try:
-            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
-            )
-            self._tng_dataloader = loader
-        except Exception as e:
-            raise e
-            
-    return self._tng_dataloader
+def on_load_checkpoint(self, checkpoint):
+    # 99% of the time you don't need to implement this method
+    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']
 ```
 
 --- 
 ### val_dataloader 
 
 ``` {.python}
-@property
-def tng_dataloader(self)
+@pl.data_loader
+def val_dataloader(self)
 ```
-Called by lightning during validation loop. Define it as a property.
+**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**
 
 ``` {.python}
-@property
+@pl.data_loader
 def val_dataloader(self):
-    if self._val_dataloader is None:
-        try:
-            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-            dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
-            loader = torch.utils.data.DataLoader(
-                dataset=dataset,
-                batch_size=self.hparams.batch_size,
-                shuffle=True
-            )
-            self._val_dataloader = loader
-        except Exception as e:
-            raise e
-            
-    return self._val_dataloader
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+    dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
+    loader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=self.hparams.batch_size,
+        shuffle=True
+    )
+    
+    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 
 
 ``` {.python}
-@property
+@pl.data_loader
 def test_dataloader(self)
 ```
-Called by lightning during test loop. Define it as a property.
+**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
+PyTorch DataLoader
 
 **Example**
 
 ``` {.python}
-@property
+@pl.data_loader
 def test_dataloader(self):
-    if self._test_dataloader is None:
-        try:
-            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-            dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
-            loader = torch.utils.data.DataLoader(
-                dataset=dataset,
-                batch_size=self.hparams.batch_size,
-                shuffle=True
-            )
-            self._test_dataloader = loader
-        except Exception as e:
-            raise e
-            
-    return self._test_dataloader
-```
-
---- 
-### 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
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+    dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
+    loader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=self.hparams.batch_size,
+        shuffle=True
+    )
+    
+    return loader
 ```
 
 --- 
@@ -374,7 +687,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)
@@ -391,4 +704,4 @@ def add_model_specific_args(parent_parser, root_dir):
     parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
     parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
     return parser
-```
+```

docs/LightningModule/methods.md

@@ -21,7 +21,8 @@ pretrained_model = MyLightningModule.load_from_metrics(
     map_location=None
 )
     
-# predict 
+# predict
+pretrained_model.eval()
 pretrained_model.freeze()
 y_hat = pretrained_model(x)
 ```
@@ -30,7 +31,7 @@ y_hat = pretrained_model(x)
 
 | Param  | description  |
 |---|---|
-|  weights_path | Path to a pytorch checkpoint  |
+|  weights_path | Path to a PyTorch checkpoint  |
 |  tags_csv | Path to meta_tags.csv file generated by the test-tube Experiment  |
 |  on_gpu | if True, puts model on GPU. Make sure to use transforms option if model devices have changed  |
 |  map_location | A dictionary mapping saved weight GPU devices to new GPU devices |

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    
 
 ---
@@ -38,3 +48,17 @@ self.trainer.current_epoch
 ...   
 ```   
 
+## Debugging   
+The LightningModule also offers these tricks to help debug.   
+
+---   
+#### example_input_array    
+In the LightningModule init, you can set a dummy tensor for this property
+to get a print out of sizes coming into and out of every layer.   
+```python
+def __init__(self):
+    # put the dimensions of the first input to your system
+    self.example_input_array = torch.rand(5, 28 * 28)
+```    
+
+

docs/Trainer/Checkpointing.md

@@ -2,21 +2,76 @@ 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)
 ```
 
+---
+### Restoring training session 
+You might want to not only load a model but also continue training it. Use this method to
+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.   
+``` {.python}
+from test_tube import Experiment
 
+exp = Experiment(version=a_previous_version_with_a_saved_checkpoint)
+trainer = Trainer(experiment=exp)
+
+# this fit call loads model weights and trainer state
+# the trainer continues seamlessly from where you left off
+# without having to do anything else.
+trainer.fit(model)
+```
+
+The trainer restores:  
+
+- global_step    
+- current_epoch    
+- All optimizers    
+- All lr_schedulers    
+- Model weights
+
+You can even change the logic of your model as long as the weights and "architecture" of 
+the system isn't different. If you add a layer, for instance, it might not work.   
+
+At a rough level, here's [what happens inside Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/root_module/model_saving.py#L63):   
+```python
+
+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)
+
+# uses the model you passed into trainer        
+model.load_state_dict(checkpoint['state_dict'])
+```    

docs/Trainer/Distributed training.md

@@ -3,6 +3,84 @@ Lightning makes multi-gpu training and 16 bit training trivial.
 *Note:*   
 None of the flags below require changing anything about your lightningModel definition. 
 
+---
+#### Choosing a backend  
+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 (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 (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.
+We recommend you use DistributedDataparallel even for single-node multi-GPU training. It is MUCH faster than DP but *may*
+have configuration issues depending on your cluster.
+
+For a deeper understanding of what lightning is doing, feel free to read [this guide](https://medium.com/@_willfalcon/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565).   
+
+---
+#### Distributed and 16-bit precision.    
+Due to an issue with apex and DistributedDataParallel (PyTorch and NVIDIA issue), Lightning does
+not allow 16-bit and DP training. We tried to get this to work, but it's an issue on their end.   
+
+Below are the possible configurations we support.    
+
+| 1 GPU  | 1+ GPUs  | DP  | DDP  | 16-bit  | command |
+|---|---|---|---|---|---|
+| 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)
+Trainer(gpus=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')
+```
+
+---
+#### CUDA flags   
+CUDA flags make certain GPUs visible to your script. 
+Lightning sets these for you automatically, there's NO NEED to do this yourself.
+```python
+# lightning will set according to what you give the trainer
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# 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.    
@@ -10,6 +88,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" ./
 ```
 
@@ -23,12 +116,8 @@ trainer = Trainer(amp_level='O2', use_amp=False)
 #### Single-gpu
 Make sure you're on a GPU machine. 
 ```python
-# set these flags
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-
 # DEFAULT
-trainer = Trainer(gpus=[0])
+trainer = Trainer(gpus=1)
 ```
 
 ---
@@ -36,17 +125,128 @@ 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
-# set these flags
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
+# to use DataParallel
+trainer = Trainer(gpus=8, distributed_backend='dp')
 
-# DEFAULT
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
+# RECOMMENDED use DistributedDataParallel
+trainer = Trainer(gpus=8, distributed_backend='ddp')
 ```
 
 ---
 #### Multi-node
-COMING SOON.
+Multi-node training is easily done by specifying these flags. 
+```python
+# train on 12*8 GPUs
+trainer = Trainer(gpus=8, nb_gpu_nodes=12, distributed_backend='ddp')
+```
+
+You must configure your job submission script correctly for the trainer to work. Here is an example
+script for the above trainer configuration.   
+
+```sh
+#!/bin/bash -l
+
+# SLURM SUBMIT SCRIPT
+#SBATCH --nodes=12
+#SBATCH --gres=gpu:8
+#SBATCH --ntasks-per-node=8
+#SBATCH --mem=0
+#SBATCH --time=0-02:00:00
+
+# activate conda env
+conda activate my_env
+
+# -------------------------
+# OPTIONAL
+# -------------------------
+# debugging flags (optional)
+# export NCCL_DEBUG=INFO
+# export PYTHONFAULTHANDLER=1
+
+# 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
+
+# on your cluster you might need these:
+# set the network interface
+# export NCCL_SOCKET_IFNAME=^docker0,lo
+# -------------------------
+
+# 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).   
+
+```python
+# ie: this:
+dataset = myDataset()
+dataloader = Dataloader(dataset)
+
+# becomes:
+dataset = myDataset()
+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/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

docs/Trainer/Logging.md

@@ -1,11 +1,117 @@
-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/).
+
+```{.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)
+```
+
+---
+#### 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)
+```
+
+---
+#### 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
+self.logger
+
+# add an image if using TestTubeLogger
+self.logger.experiment.add_image(...)
+```
+
+
 #### Display metrics in progress bar 
 ``` {.python}
 # DEFAULT
-trainer = Trainer(progress_bar=True)
+trainer = Trainer(show_progress_bar=True)
 ```
 
 ---
@@ -13,7 +119,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,27 +150,23 @@ 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
+
+    ...
+
+args = parser.parse_args()
+model = MyModel(args)
+
+logger = TestTubeLogger(...)
+t = Trainer(logger=logger)
+trainer.fit(model)
 ```
 
----
-#### 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)
-```
-
-
 ---
 #### Write logs file to csv every k batches 
 Every k batches, lightning will write the new logs to disk

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

@@ -1,4 +1,4 @@
-The lightning training loop handles everything except the actual computations of your model. To decide what will happen in your training loop, define the [training_step function](../../Pytorch-lightning/LightningModule/#training_step).
+The lightning training loop handles everything except the actual computations of your model. To decide what will happen in your training loop, define the [training_step function](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#training_step).
 
 Below are all the things lightning automates for you in the training loop.
 
@@ -11,17 +11,6 @@ Accumulated gradients runs K small batches of size N before doing a backwards pa
 trainer = Trainer(accumulate_grad_batches=1)
 ```
 
----
-#### Anneal Learning rate
-Cut the learning rate by 10 at every epoch listed in this list.
-``` {.python}
-# DEFAULT (don't anneal)
-trainer = Trainer(lr_scheduler_milestones=None)
-
-# cut LR by 10 at 100, 200, and 300 epochs 
-trainer = Trainer(lr_scheduler_milestones='100, 200, 300')
-```
-
 ---
 #### Force training for min or max epochs
 It can be useful to force training for a minimum number of epochs or limit to a max number
@@ -30,6 +19,25 @@ It can be useful to force training for a minimum number of epochs or limit to a
 trainer = Trainer(min_nb_epochs=1, max_nb_epochs=1000)
 ```
 
+---
+#### Early stopping
+The trainer already sets up default early stopping for you. 
+To modify this behavior, pass in your own EarlyStopping callback.
+``` {.python}
+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'
+)
+
+trainer = Trainer(early_stop_callback=early_stop_callback)
+```
+
 ---
 #### 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)
@@ -39,15 +47,18 @@ trainer = Trainer(enable_early_stop=True)
 ```
 
 ---
-#### Gradient Clipping 
-Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
+#### 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=0)
+trainer = Trainer(gradient_clip_val=0)
+
+# clip gradients with norm above 0.5
+trainer = Trainer(gradient_clip_val=0.5)
 ```
 
-
-
 ---
 #### Inspect gradient norms
 Looking at grad norms can help you figure out where training might be going wrong.
@@ -62,7 +73,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)

docs/Trainer/Validation loop.md

@@ -1,12 +1,10 @@
-The lightning validation loop handles everything except the actual computations of your model. To decide what will happen in your validation loop, define the [validation_step function](../../Pytorch-lightning/LightningModule/#validation_step).
+The lightning validation loop handles everything except the actual computations of your model. To decide what will happen in your validation loop, define the [validation_step function](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/#validation_step).
 Below are all the things lightning automates for you in the validation loop.
 
 **Note**   
 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)
@@ -54,4 +58,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)

docs/Trainer/hooks.md

@@ -0,0 +1,131 @@
+# Hooks
+[[Github Code](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/root_module/hooks.py)]   
+
+There are cases when you might want to do something different at different parts of the training/validation loop.
+To enable a hook, simply override the method in your LightningModule and the trainer will call it at the correct time.
+
+**Contributing** If there's a hook you'd like to add, simply:    
+1. Fork PyTorchLightning.    
+2. Add the hook [here](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/root_module/hooks.py).       
+3. Add the correct place in the [Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/trainer.py) where it should be called.    
+
+---
+#### on_epoch_start
+Called in the training loop at the very beginning of the epoch.   
+```python
+def on_epoch_start(self):
+    # do something when the epoch starts
+```
+
+---
+#### on_epoch_end
+Called in the training loop at the very end of the epoch.   
+```python
+def on_epoch_end(self):
+    # do something when the epoch ends 
+```
+
+---
+#### on_batch_start
+Called in the training loop before anything happens for that batch.   
+```python
+def on_batch_start(self):
+    # do something when the batch starts
+```
+
+---
+#### on_batch_end
+Called in the training loop after the batch.   
+```python
+def on_batch_end(self):
+    # do something when the batch ends 
+```
+
+---
+#### on_pre_performance_check
+Called at the very beginning of the validation loop.   
+```python
+def on_pre_performance_check(self):
+    # do something before validation starts 
+```
+
+---
+#### on_post_performance_check
+Called at the very end of the validation loop.   
+```python
+def on_post_performance_check(self):
+    # do something before validation end
+```
+
+---
+#### 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
+# 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.
+Good place to inspect weight information with weights updated.
+
+Called once per optimizer
+```python
+def on_before_zero_grad(self, optimizer):
+    # do something with the optimizer or inspect it. 
+```
+
+---
+#### on_after_backward
+Called in the training loop after model.backward()
+This is the ideal place to inspect or log gradient information 
+```python
+def on_after_backward(self):
+    # example to inspect gradient information in tensorboard
+    if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
+        params = self.state_dict()
+        for k, v in params.items():
+            grads = v
+            name = k
+            self.logger.experiment.add_histogram(tag=name, values=grads, global_step=self.trainer.global_step)
+```

docs/Trainer/index.md

@@ -19,56 +19,69 @@ But of course the fun is in all the advanced things it can do:
 
 **Checkpointing**    
 
-- Model saving
-- Model loading 
+- [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)**    
 
-- [Running grid search on a cluster](SLURM%20Managed%20Cluster/#running-grid-search-on-a-cluster)
-- [Walltime auto-resubmit](SLURM%20Managed%20Cluster/#walltime-auto-resubmit)   
+- [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**  
 
-- [Fast dev run](Debugging/#fast-dev-run)
-- [Inspect gradient norms](Debugging/#inspect-gradient-norms)
-- [Log GPU usage](Debugging/#Log-gpu-usage)
-- [Make model overfit on subset of data](Debugging/#make-model-overfit-on-subset-of-data)
-- [Print the parameter count by layer](Debugging/#print-the-parameter-count-by-layer)
-- [Pring which gradients are nan](Debugging/#print-which-gradients-are-nan)
+- [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)
+- [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)
+- [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**    
 
-- [16-bit mixed precision](Distributed%20training/#16-bit-mixed-precision)
-- [Multi-GPU](Distributed%20training/#Multi-GPU)
-- [Multi-node](Distributed%20training/#Multi-node)
-- [Single GPU](Distributed%20training/#single-gpu)
-- [Self-balancing architecture](Distributed%20training/#self-balancing-architecture)
+- [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)
+- [Single GPU](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#single-gpu)
+- [Self-balancing architecture](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/#self-balancing-architecture)
 
 
 **Experiment Logging**   
 
-- [Display metrics in progress bar](Logging/#display-metrics-in-progress-bar)
-- Log arbitrary metrics
-- [Log metric row every k batches](Logging/#log-metric-row-every-k-batches)
-- [Process position](Logging/#process-position)
-- [Save a snapshot of all hyperparameters](Logging/#save-a-snapshot-of-all-hyperparameters) 
-- [Snapshot code for a training run](Logging/#snapshot-code-for-a-training-run) 
-- [Write logs file to csv every k batches](Logging/#write-logs-file-to-csv-every-k-batches)
+- [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)
+- [Process position](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#process-position)
+- [Tensorboard support](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#tensorboard-support)
+- [Save a snapshot of all hyperparameters](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#save-a-snapshot-of-all-hyperparameters) 
+- [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**    
 
-- [Accumulate gradients](Training%20Loop/#accumulated-gradients)
-- [Anneal Learning rate](Training%20Loop/#anneal-learning-rate)
-- [Force training for min or max epochs](Training%20Loop/#force-training-for-min-or-max-epochs)
-- [Force disable early stop](Training%20Loop/#force-disable-early-stop)
-- [Use multiple optimizers (like GANs)](../Pytorch-lightning/LightningModule/#configure_optimizers)
-- [Set how much of the training set to check (1-100%)](Training%20Loop/#set-how-much-of-the-training-set-to-check)
+- [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**    
 
-- [Check validation every n epochs](Validation%20Loop/#check-validation-every-n-epochs)
-- [Set how much of the validation set to check](Validation%20Loop/#set-how-much-of-the-validation-set-to-check)
-- [Set how much of the test set to check](Validation%20Loop/#set-how-much-of-the-test-set-to-check)
-- [Set validation check frequency within 1 training epoch](Validation%20Loop/#set-validation-check-frequency-within-1-training-epoch)
-- [Set the number of validation sanity steps](Validation%20Loop/#set-the-number-of-validation-sanity-steps)
+- [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/)
+- [Set how much of the validation set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-validation-set-to-check)
+- [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

@@ -3,7 +3,7 @@ In 99% of cases you want to just copy [this template](https://github.com/william
 
 ```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/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.

docs/index.md

@@ -1,10 +1,67 @@
 ###### New project Quick Start    
-To start a new project define these two files.   
+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.    
 
-1. [Define a LightningModule](/LightningModule/RequiredTrainerInterface/#template-model-definition)  
-2. Pick a trainer      
-    - [Basic CPU Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py) 
-    - [GPU cluster Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_gpu_cluster_template.py)
+###### 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
+```   
+
+###### 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, GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/tree/master/examples/basic_examples) 
+    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/tree/master/examples/multi_node_examples)
 
 ###### Docs shortcuts
 - [LightningModule](LightningModule/RequiredTrainerInterface/)  
@@ -19,8 +76,10 @@ To start a new project define these two files.
 
 ###### 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)    
 
@@ -35,6 +94,7 @@ To start a new project define these two files.
 - [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 input and output size of every module in system](https://williamfalcon.github.io/pytorch-lightning/LightningModule/properties/#example_input_array)
 
 
 ###### Distributed training    
@@ -49,9 +109,9 @@ To start a new project define these two files.
 ###### Experiment Logging   
 
 - [Display metrics in progress bar](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#display-metrics-in-progress-bar)
-- Log arbitrary metrics
 - [Log metric row every k batches](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#log-metric-row-every-k-batches)
 - [Process position](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#process-position)
+- [Tensorboard support](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#tensorboard-support)
 - [Save a snapshot of all hyperparameters](https://williamfalcon.github.io/pytorch-lightning/Trainer/Logging/#save-a-snapshot-of-all-hyperparameters) 
 - [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)
@@ -59,18 +119,24 @@ To start a new project define these two files.
 ###### Training loop    
 
 - [Accumulate gradients](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#accumulated-gradients)
-- [Anneal Learning rate](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#anneal-learning-rate)
 - [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)
-- [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/Pytorch-Lightning/LightningModule/#configure_optimizers)
+- [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/)
 - [Set how much of the validation set to check](https://williamfalcon.github.io/pytorch-lightning/Trainer/Validation%20loop/#set-how-much-of-the-validation-set-to-check)
 - [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/requirements.txt

@@ -0,0 +1,2 @@
+mkdocs-material==4.4.0
+mkdocs==1.0.4

docs/source/_static/coverage.svg

@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" width="99" height="20">
+    <linearGradient id="b" x2="0" y2="100%">
+        <stop offset="0" stop-color="#bbb" stop-opacity=".1"/>
+        <stop offset="1" stop-opacity=".1"/>
+    </linearGradient>
+    <mask id="a">
+        <rect width="99" height="20" rx="3" fill="#fff"/>
+    </mask>
+    <g mask="url(#a)">
+        <path fill="#555" d="M0 0h63v20H0z"/>
+        <path fill="#4c1" d="M63 0h36v20H63z"/>
+        <path fill="url(#b)" d="M0 0h99v20H0z"/>
+    </g>
+    <g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11">
+        <text x="31.5" y="15" fill="#010101" fill-opacity=".3">coverage</text>
+        <text x="31.5" y="14">coverage</text>
+        <text x="80" y="15" fill="#010101" fill-opacity=".3">99%</text>
+        <text x="80" y="14">99%</text>
+    </g>
+</svg>

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.

examples/__init__.py

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

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
+```

examples/basic_examples/cpu_template.py

@@ -0,0 +1,53 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import numpy as np
+import torch
+
+from argparse import ArgumentParser
+from pytorch_lightning import Trainer
+from examples.basic_examples.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:
+    """
+    # ------------------------
+    # 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)

examples/basic_examples/gpu_template.py

@@ -0,0 +1,78 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import numpy as np
+import torch
+
+from argparse import ArgumentParser
+from pytorch_lightning import Trainer
+from examples.basic_examples.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:
+    """
+    # ------------------------
+    # 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)

examples/basic_examples/lightning_module_template.py

@@ -1,3 +1,6 @@
+"""
+Example template for defining a system
+"""
 import os
 from collections import OrderedDict
 import torch.nn as nn
@@ -5,11 +8,12 @@ from torchvision.datasets import MNIST
 import torchvision.transforms as transforms
 import torch
 import torch.nn.functional as F
-from test_tube import HyperOptArgumentParser
+from argparse import ArgumentParser
 from torch import optim
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
 
+import pytorch_lightning as pl
 from pytorch_lightning.root_module.root_module import LightningModule
 
 
@@ -24,10 +28,14 @@ class LightningTemplateModel(LightningModule):
         :param hparams:
         """
         # init superclass
-        super(LightningTemplateModel, self).__init__(hparams)
+        super(LightningTemplateModel, 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)
+
         # build model
         self.__build_model()
 
@@ -39,11 +47,13 @@ class LightningTemplateModel(LightningModule):
         Layout model
         :return:
         """
-        self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
+        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)
+        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim,
+                              out_features=self.hparams.out_features)
 
     # ---------------------
     # TRAINING
@@ -69,20 +79,25 @@ 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)
 
         # 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 or self.trainer.use_ddp2:
+            loss_val = loss_val.unsqueeze(0)
+
         output = OrderedDict({
             'loss': loss_val
         })
@@ -90,13 +105,13 @@ class LightningTemplateModel(LightningModule):
         # 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)
 
@@ -105,10 +120,19 @@ class LightningTemplateModel(LightningModule):
         # 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 or self.trainer.use_ddp2:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
 
         output = OrderedDict({
             'val_loss': loss_val,
-            'val_acc': torch.tensor(val_acc).cuda(loss_val.device.index),
+            'val_acc': val_acc,
         })
 
         # can also return just a scalar instead of a dict (return loss_val)
@@ -127,27 +151,25 @@ class LightningTemplateModel(LightningModule):
         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 = 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 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_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
-        return tqdm_dic
-
-    def update_tng_log_metrics(self, logs):
-        return logs
-
-    # ---------------------
-    # MODEL SAVING
-    # ---------------------
-    def get_save_dict(self):
-        checkpoint = {'state_dict': self.state_dict()}
-        return checkpoint
-
-    def load_model_specific(self, checkpoint):
-        self.load_state_dict(checkpoint['state_dict'])
-        pass
+        tqdm_dict = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean}
+        result = {'progress_bar': tqdm_dict, 'logs': tqdm_dict}
+        return result
 
     # ---------------------
     # TRAINING SETUP
@@ -158,92 +180,74 @@ class LightningTemplateModel(LightningModule):
         :return: list of optimizers
         """
         optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
-        return [optimizer]
+        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
+        return [optimizer], [scheduler]
 
     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)
+        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
+        # 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 as e:
-            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
 
-    @property
-    def tng_dataloader(self):
-        if self._tng_dataloader is None:
-            try:
-                self._tng_dataloader = self.__dataloader(train=True)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._tng_dataloader
+    @pl.data_loader
+    def train_dataloader(self):
+        print('training data loader called')
+        return self.__dataloader(train=True)
 
-    @property
+    @pl.data_loader
     def val_dataloader(self):
-        if self._val_dataloader is None:
-            try:
-                self._val_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._val_dataloader
+        print('val data loader called')
+        return self.__dataloader(train=False)
 
-    @property
+    @pl.data_loader
     def test_dataloader(self):
-        if self._test_dataloader is None:
-            try:
-                self._test_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._test_dataloader
+        print('test data loader called')
+        return self.__dataloader(train=False)
 
     @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):
+    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])
+        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.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('--in_features', default=28 * 28, type=int)
         parser.add_argument('--out_features', default=10, type=int)
-        parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference
+        # use 500 for CPU, 50000 for GPU to see speed difference
+        parser.add_argument('--hidden_dim', default=50000, type=int)
+        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 the 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

examples/domain_templates/gan.py

@@ -0,0 +1,178 @@
+"""
+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
+"""
+from argparse import ArgumentParser
+import os
+import numpy as np
+
+import torchvision
+import torchvision.transforms as transforms
+from torchvision.datasets import MNIST
+
+from torch.utils.data import DataLoader
+
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+import pytorch_lightning as pl
+from test_tube import Experiment
+
+
+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
+
+    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
+
+        # 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)
+
+            return g_loss
+
+        # 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
+
+            return d_loss
+
+    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 main(hparams):
+    # save tensorboard logs
+    exp = Experiment(save_dir=os.getcwd())
+
+    # init model
+    model = GAN(hparams)
+
+    # fit trainer on CPU
+    trainer = pl.Trainer(experiment=exp, max_nb_epochs=200)
+    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)

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
+```

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

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

examples/multi_node_examples/multi_node_ddp2_demo.py

@@ -0,0 +1,55 @@
+"""
+Multi-node example (GPU)
+"""
+import os
+import numpy as np
+import torch
+
+from argparse import ArgumentParser
+from pytorch_lightning import Trainer
+from examples.basic_examples.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
+    # ------------------------
+    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)

examples/multi_node_examples/multi_node_ddp_demo.py

@@ -0,0 +1,55 @@
+"""
+Multi-node example (GPU)
+"""
+import os
+import numpy as np
+import torch
+
+from argparse import ArgumentParser
+from pytorch_lightning import Trainer
+from examples.basic_examples.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
+    # ------------------------
+    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)

examples/new_project_templates/__init__.py

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

examples/new_project_templates/multi_node_cluster_template.py

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

examples/new_project_templates/single_cpu_template.py

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

examples/new_project_templates/single_gpu_node_16bit_template.py

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

examples/new_project_templates/single_gpu_node_dp_template.py

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

examples/new_project_templates/single_gpu_node_template.py

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

examples/new_project_templates/single_gpu_template.py

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

examples/new_project_templates/trainer_cpu_template.py

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

mkdocs.yml

@@ -1,10 +1,16 @@
-site_name: Pytorch lightning Documentation
+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.'
+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

pytorch_lightning/__init__.py

@@ -1,2 +1,9 @@
-from .models import Trainer
-from .root_module.root_module import LightningModule
+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 +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,5 +1,9 @@
+import os
+import shutil
+import warnings
+
 import numpy as np
-import os, shutil
+
 from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
 
 
@@ -8,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
@@ -119,9 +122,9 @@ 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
-            )
+            print('Early stopping conditioned on metric `%s` '
+                  'which is not available. Available metrics are: %s' %
+                  (self.monitor, ','.join(list(logs.keys()))), RuntimeWarning)
             exit(-1)
 
         if self.monitor_op(current - self.min_delta, self.best):
@@ -185,8 +188,7 @@ class ModelCheckpoint(Callback):
 
         if mode not in ['auto', 'min', 'max']:
             print('ModelCheckpoint mode %s is unknown, '
-                          'fallback to auto mode.' % (mode),
-                          RuntimeWarning)
+                  'fallback to auto mode.' % (mode), RuntimeWarning)
             mode = 'auto'
 
         if mode == 'min':
@@ -230,8 +232,8 @@ 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)
+                    print('Can save best model only with %s available,'
+                          ' skipping.' % (self.monitor), RuntimeWarning)
                 else:
                     if self.monitor_op(current, self.best):
                         if self.verbose > 0:
@@ -252,6 +254,37 @@ class ModelCheckpoint(Callback):
                 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]
@@ -260,4 +293,3 @@ if __name__ == '__main__':
         print(loss)
         if should_stop:
             break
-

pytorch_lightning/logging/__init__.py

@@ -0,0 +1,10 @@
+from .base import LightningLoggerBase, rank_zero_only
+
+try:
+    from .test_tube_logger import TestTubeLogger
+except ModuleNotFoundError:
+    pass
+try:
+    from .mlflow_logger import MLFlowLogger
+except ModuleNotFoundError:
+    pass

pytorch_lightning/logging/base.py

@@ -0,0 +1,71 @@
+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 version(self):
+        """Return the experiment version"""
+        return None

pytorch_lightning/logging/mlflow_logger.py

@@ -0,0 +1,56 @@
+from time import time
+from logging import getLogger
+
+import mlflow
+
+from .base import LightningLoggerBase, rank_zero_only
+
+logger = getLogger(__name__)
+
+
+class MLFlowLogger(LightningLoggerBase):
+    def __init__(self, experiment_name, tracking_uri=None):
+        super().__init__()
+        self.client = mlflow.tracking.MlflowClient(tracking_uri)
+        self.experiment_name = experiment_name
+        self._run_id = None
+
+    @property
+    def run_id(self):
+        if self._run_id is not None:
+            return self._run_id
+
+        experiment = self.client.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.client.create_experiment(self.experiment_name)
+            experiment = self.client.get_experiment_by_name(self.experiment_name)
+
+        run = self.client.create_run(experiment.experiment_id)
+        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.client.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.client.log_metric(self.run_id, k, v, timestamp_ms, step_num)
+
+    def save(self):
+        pass
+
+    @rank_zero_only
+    def finalize(self, status="FINISHED"):
+        self.client.set_terminated(self.run_id, status)

pytorch_lightning/logging/test_tube_logger.py

@@ -0,0 +1,93 @@
+import os.path
+from copy import copy
+
+from .base import LightningLoggerBase, rank_zero_only
+
+from test_tube import Experiment
+
+
+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):
+        self.experiment.argparse(params)
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step_num=None):
+        self.experiment.log(metrics, global_step=step_num)
+
+    @rank_zero_only
+    def save(self):
+        self.experiment.save()
+
+    @rank_zero_only
+    def finalize(self, status):
+        self.save()
+        self.close()
+
+    def close(self):
+        self.experiment.close()
+
+    @property
+    def rank(self):
+        if self._experiment is None:
+            return self._rank
+        else:
+            return self.experiment.rank
+
+    @rank.setter
+    def rank(self, value):
+        if self._experiment is None:
+            self._rank = value
+        else:
+            return self.experiment.rank
+
+    @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/models/__init__.py

@@ -1 +0,0 @@
-from .trainer import Trainer

pytorch_lightning/models/sample_model_template/model_template.py

@@ -1,203 +0,0 @@
-import torch.nn as nn
-import numpy as np
-from pytorch_lightning.root_module.root_module import LightningModule
-from test_tube import HyperOptArgumentParser
-from torchvision.datasets import MNIST
-import torchvision.transforms as transforms
-import torch
-import torch.nn.functional as F
-
-
-class ExampleModel1(LightningModule):
-    """
-    Sample model to show how to define a template
-    """
-
-    def __init__(self, hparams):
-        # init superclass
-        super(ExampleModel1, self).__init__(hparams)
-
-        self.batch_size = hparams.batch_size
-
-        # build model
-        self.__build_model()
-
-    # ---------------------
-    # MODEL SETUP
-    # ---------------------
-    def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
-        self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
-        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
-        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
-
-        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features)
-
-    # ---------------------
-    # TRAINING
-    # ---------------------
-    def forward(self, x):
-        x = self.c_d1(x)
-        x = F.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):
-        """
-        Called inside the training loop
-        :param data_batch:
-        :return:
-        """
-        # forward pass
-        x, y = data_batch
-        x = x.view(x.size(0), -1)
-        y_hat = self.forward(x)
-
-        # calculate loss
-        loss_val = self.loss(y, y_hat)
-
-        tqdm_dic = {'jefe': 1}
-        return loss_val, tqdm_dic
-
-    def validation_step(self, data_batch):
-        """
-        Called inside the validation loop
-        :param data_batch:
-        :return:
-        """
-        x, y = data_batch
-        x = x.view(x.size(0), -1)
-        y_hat = self.forward(x)
-
-        loss_val = self.loss(y, y_hat)
-
-        # acc
-        labels_hat = torch.argmax(y_hat, dim=1)
-        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-
-        output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
-        return output
-
-    def validation_end(self, outputs):
-        """
-        Called at the end of validation to aggregate outputs
-        :param outputs: list of individual outputs of each validation step
-        :return:
-        """
-        val_loss_mean = 0
-        accs = []
-        for output in outputs:
-            val_loss_mean += output['val_loss']
-            accs.append(output['val_acc'])
-
-        val_loss_mean /= len(outputs)
-        tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)}
-        return tqdm_dic
-
-    def update_tng_log_metrics(self, logs):
-        return logs
-
-    # ---------------------
-    # MODEL SAVING
-    # ---------------------
-    def get_save_dict(self):
-        checkpoint = {
-            'state_dict': self.state_dict(),
-        }
-
-        return checkpoint
-
-    def load_model_specific(self, checkpoint):
-        self.load_state_dict(checkpoint['state_dict'])
-        pass
-
-    # ---------------------
-    # TRAINING SETUP
-    # ---------------------
-    def configure_optimizers(self):
-        """
-        return whatever optimizers we want here
-        :return: list of optimizers
-        """
-        optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
-        self.optimizers = [optimizer]
-        return self.optimizers
-
-    def __dataloader(self, train):
-        # init data generators
-        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-
-        dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
-
-        loader = torch.utils.data.DataLoader(
-            dataset=dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True
-        )
-
-        return loader
-
-    @property
-    def tng_dataloader(self):
-        if self._tng_dataloader is None:
-            try:
-                self._tng_dataloader = self.__dataloader(train=True)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._tng_dataloader
-
-    @property
-    def val_dataloader(self):
-        if self._val_dataloader is None:
-            try:
-                self._val_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._val_dataloader
-
-    @property
-    def test_dataloader(self):
-        if self._test_dataloader is None:
-            try:
-                self._test_dataloader = self.__dataloader(train=False)
-            except Exception as e:
-                print(e)
-                raise e
-        return self._test_dataloader
-
-    @staticmethod
-    def add_model_specific_args(parent_parser):
-        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
-
-        # param overwrites
-        # parser.set_defaults(gradient_clip=5.0)
-
-        # network params
-        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
-        parser.add_argument('--in_features', default=28*28)
-        parser.add_argument('--hidden_dim', default=500)
-        parser.add_argument('--out_features', default=10)
-
-        # data
-        parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str)
-
-        # training params (opt)
-        parser.opt_list('--learning_rate', default=0.001, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
-        parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
-        parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
-        return parser

pytorch_lightning/models/trainer.py

@@ -1,831 +0,0 @@
-"""
-The trainer handles all the logic for running a val loop, training loop, distributing, etc...
-"""
-from time import sleep
-import subprocess
-import traceback
-import warnings
-import os
-import pdb
-
-import torch
-from torch.utils.data.distributed import DistributedSampler
-from torch.optim.lr_scheduler import MultiStepLR
-import torch.multiprocessing as mp
-import torch.distributed as dist
-import numpy as np
-import tqdm
-
-from pytorch_lightning.root_module.memory import get_gpu_memory_map
-from pytorch_lightning.root_module.model_saving import TrainerIO
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
-
-
-try:
-    from apex import amp
-    APEX_AVAILABLE = True
-except ModuleNotFoundError:
-    APEX_AVAILABLE = False
-
-
-def reduce_distributed_output(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] = reduce_distributed_output(output[k], nb_gpus)
-
-        # reduce only metrics that have the same nb of gpus
-        elif output[k].size(0) == nb_gpus:
-            reduced = torch.mean(output[k])
-            output[k] = reduced
-    return output
-
-
-class Trainer(TrainerIO):
-
-    def __init__(self,
-                 experiment,
-                 early_stop_callback=None,
-                 checkpoint_callback=None,
-                 gradient_clip=0,
-                 cluster=None,
-                 process_position=0,
-                 current_gpu_name=0,
-                 nb_gpu_nodes=1,
-                 gpus=None,
-                 progress_bar=True,
-                 overfit_pct=0.0,
-                 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=0.95,
-                 log_save_interval=100, add_log_row_interval=10,
-                 lr_scheduler_milestones=None,
-                 distributed_backend='dp',
-                 use_amp=False,
-                 print_nan_grads=False,
-                 print_weights_summary=True,
-                 amp_level='O2',
-                 nb_sanity_val_steps=5):
-
-        """
-
-        :param experiment: Test-tube experiment
-        :param early_stop_callback: from pytorch_lightning import EarlyStopping
-        :param checkpoint_callback: from pytorch_lightning import Checkpoint
-        :param gradient_clip:
-        :param cluster:
-        :param process_position:
-        :param current_gpu_name:
-        :param nb_gpu_nodes:
-        :param gpus:
-        :param progress_bar:
-        :param overfit_pct:
-        :param track_grad_norm:
-        :param check_val_every_n_epoch:
-        :param fast_dev_run:
-        :param accumulate_grad_batches:
-        :param max_nb_epochs:
-        :param min_nb_epochs:
-        :param train_percent_check:
-        :param val_percent_check:
-        :param test_percent_check:
-        :param val_check_interval:
-        :param log_save_interval:
-        :param add_log_row_interval:
-        :param lr_scheduler_milestones:
-        :param distributed_backend: 'np' to use DistributedParallel, 'ddp' to use DistributedDataParallel
-        :param use_amp:
-        :param print_nan_grads:
-        :param print_weights_summary:
-        :param amp_level:
-        :param nb_sanity_val_steps:
-        """
-
-        # Transfer params
-        self.nb_gpu_nodes = nb_gpu_nodes
-        self.gradient_clip = gradient_clip
-        self.check_val_every_n_epoch = check_val_every_n_epoch
-        self.enable_early_stop = early_stop_callback is not None
-        self.track_grad_norm = track_grad_norm
-        self.fast_dev_run = fast_dev_run
-        self.on_gpu = gpus is not None and torch.cuda.is_available()
-        self.progress_bar = progress_bar
-        self.experiment = experiment
-        self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
-        self.cluster = cluster
-        self.process_position = process_position
-        self.current_gpu_name = current_gpu_name
-        self.print_weights_summary = print_weights_summary
-        self.checkpoint_callback = checkpoint_callback
-
-        if self.checkpoint_callback is not None:
-            self.checkpoint_callback.save_function = self.save_checkpoint
-
-        self.early_stop = early_stop_callback
-        self.model = None
-        self.max_nb_epochs = max_nb_epochs
-        self.accumulate_grad_batches = accumulate_grad_batches
-        self.early_stop_callback = early_stop_callback
-        self.min_nb_epochs = min_nb_epochs
-        self.nb_sanity_val_steps = nb_sanity_val_steps
-        self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')]
-        self.lr_schedulers = []
-        self.amp_level = amp_level
-        self.print_nan_grads = print_nan_grads
-        self.data_parallel_device_ids = None
-        self.world_size = 1
-        self.use_ddp = False
-        self.use_dp = False
-
-        # gpus come in as a string.
-        # if gpus = -1 then use all available devices
-        # otherwise, split the string using commas
-        if gpus is not None:
-            if gpus == '-1':
-                self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
-            else:
-                self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
-
-            # set the correct cuda visible devices (using pci order)
-            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-            os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
-            print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
-
-        # make DP and DDP mutually exclusive
-        # single GPU will also use DP with devices=[0]
-        have_gpus = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
-        if have_gpus:
-            self.use_dp = distributed_backend == 'dp'
-            self.use_ddp = distributed_backend == 'ddp'
-
-        # process info
-        self.proc_rank = 0
-
-        # training state
-        self.optimizers = None
-        self.prog_bar = None
-        self.global_step = 0
-        self.current_epoch = 0
-        self.total_batches = 0
-
-        # logging
-        self.log_save_interval = log_save_interval
-        self.val_check_interval = val_check_interval
-        self.add_log_row_interval = add_log_row_interval
-
-        # dataloaders
-        self.tng_dataloader = None
-        self.test_dataloader = None
-        self.val_dataloader = None
-
-        # how much of the data to use
-        self.__determine_data_use_amount(train_percent_check, val_percent_check, test_percent_check, overfit_pct)
-        print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
-
-        # 16 bit mixed precision training using apex
-        self.use_amp = use_amp and APEX_AVAILABLE
-        if self.use_amp:
-            print('using 16bit precision')
-
-        if use_amp and not APEX_AVAILABLE:
-            msg = '''
-            You set use_amp=True but do not have apex installed.
-            Install apex first using this guide and rerun with use_amp=True: 
-            https://github.com/NVIDIA/apex#linux
-            
-            this run will NOT use 16 bit precision
-            '''
-            warnings.warn(msg)
-
-    @property
-    def data_parallel(self):
-        return self.use_dp or self.use_ddp
-
-    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
-
-    def __get_model(self):
-        return self.model.module if self.data_parallel else self.model
-
-    def __is_function_implemented(self, f_name):
-        model = self.__get_model()
-        f_op = getattr(model, f_name, None)
-        return callable(f_op)
-
-    @property
-    def __tng_tqdm_dic(self):
-        tqdm_dic = {
-            'tng_loss': '{0:.3f}'.format(self.avg_loss),
-            'v_nb': '{}'.format(self.experiment.version),
-            'epoch': '{}'.format(self.current_epoch),
-            'batch_nb':'{}'.format(self.batch_nb),
-        }
-        tqdm_dic.update(self.tqdm_metrics)
-
-        if self.on_gpu:
-            tqdm_dic['gpu'] = '{}'.format(self.current_gpu_name)
-
-        return tqdm_dic
-
-    def __layout_bookeeping(self):
-        # training bookeeping
-        self.total_batch_nb = 0
-        self.running_loss = []
-        self.avg_loss = 0
-        self.batch_nb = 0
-        self.tqdm_metrics = {}
-
-        # determine number of training batches
-        self.nb_tng_batches = len(self.tng_dataloader)
-        self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
-
-        # determine number of validation batches
-        self.nb_val_batches = len(self.val_dataloader)
-        self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
-        self.nb_val_batches = max(1, self.nb_val_batches)
-        self.nb_val_batches = self.nb_val_batches
-
-        # determine number of test batches
-        self.nb_test_batches = len(self.test_dataloader)
-        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
-
-        # determine when to check validation
-        self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval)
-
-    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 validate(self, model, dataloader, max_batches):
-        """
-        Run validation code
-        :param model: PT model
-        :param dataloader: PT dataloader
-        :param max_batches: Scalar
-        :return:
-        """
-        # enable eval mode
-        model.zero_grad()
-        model.eval()
-
-        # disable gradients to save memory
-        torch.set_grad_enabled(False)
-
-        # bookkeeping
-        outputs = []
-
-        # run training
-        for batch_i, data_batch in enumerate(dataloader):
-
-            if data_batch is None:
-                continue
-
-            # stop short when on fast dev run
-            if max_batches is not None and batch_i >= max_batches:
-                break
-
-            # -----------------
-            # RUN VALIDATION STEP
-            # -----------------
-            if self.use_ddp:
-                output = model(data_batch, batch_i)
-            elif self.use_dp:
-                output = model(data_batch, batch_i)
-                output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
-
-            else:
-                output = model.validation_step(data_batch, batch_i)
-
-            outputs.append(output)
-
-            # batch done
-            if self.progress_bar and self.prog_bar is not None:
-                self.prog_bar.update(1)
-
-        # give model a chance to do something with the outputs
-        if self.data_parallel:
-            val_results = model.module.validation_end(outputs)
-        else:
-            val_results = model.validation_end(outputs)
-
-        # enable train mode again
-        model.train()
-
-        # enable gradients to save memory
-        torch.set_grad_enabled(True)
-
-        return val_results
-
-    def __get_dataloaders(self, model):
-        """
-        Dataloaders are provided by the model
-        :param model:
-        :return:
-        """
-        self.tng_dataloader = model.tng_dataloader
-        self.test_dataloader = model.test_dataloader
-        self.val_dataloader = model.val_dataloader
-
-        if self.use_ddp and not isinstance(self.tng_dataloader.sampler, DistributedSampler):
-            msg = '''
-            when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
-            
-            ie: this:
-            dataset = myDataset()
-            dataloader = Dataloader(dataset)
-            
-            becomes:
-            dataset = myDataset()
-            dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
-            dataloader = Dataloader(dataset, sampler=dist_sampler)
-            '''
-            raise Exception(msg)
-
-    # -----------------------------
-    # MODEL TRAINING
-    # -----------------------------
-    def fit(self, model):
-
-        # when using multi-node or DDP within a node start each module in a separate process
-        if self.use_ddp:
-            # must copy only the meta of the exp so it survives pickle/unpickle when going to new process
-            self.experiment = self.experiment.get_meta_copy()
-
-            # whenever we have the correct number of tasks, we let slurm manage processes
-            # otherwise we launch the required number of processes
-            try:
-                nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
-                nb_requested_gpus = len(self.data_parallel_device_ids)
-                is_slurm_managing_tasks = nb_slurm_tasks == nb_requested_gpus
-            except Exception as e:
-                # likely not on slurm, so set the slurm managed flag to false
-                is_slurm_managing_tasks = False
-
-            if is_slurm_managing_tasks:
-                task = int(os.environ['SLURM_LOCALID'])
-                self.ddp_train(task, model)
-            else:
-                msg = f"""
-                You requested {nb_requested_gpus} GPUs but launched {nb_slurm_tasks} slurm tasks. 
-                We will launch {nb_requested_gpus} processes for you. 
-                We recommend you let slurm manage the processes by setting: --ntasks-per-node={nb_requested_gpus}
-                If you're not using SLURM, ignore this message!
-                """
-                warnings.warn(msg)
-                mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), 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)
-
-        # ON CPU
-        else:
-            # CHOOSE OPTIMIZER
-            # filter out the weights that were done on gpu so we can load on good old cpus
-            self.optimizers = model.configure_optimizers()
-
-            # run through amp wrapper
-            if self.use_amp:
-                # An example
-                model, optimizers = amp.initialize(
-                    model, self.optimizers, opt_level=self.amp_level,
-                )
-                self.optimizers = optimizers
-
-            self.__run_pretrain_routine(model)
-
-    def dp_train(self, model):
-
-        # CHOOSE OPTIMIZER
-        # filter out the weights that were done on gpu so we can load on good old cpus
-        self.optimizers = model.configure_optimizers()
-
-        model.cuda(self.data_parallel_device_ids[0])
-        model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
-
-        # run through amp wrapper
-        if self.use_amp:
-            # An example
-            model, optimizers = amp.initialize(
-                model, self.optimizers, opt_level=self.amp_level,
-            )
-            self.optimizers = optimizers
-
-        self.__run_pretrain_routine(model)
-
-    def 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_rank = int(os.environ['SLURM_NODEID'])
-        except KeyError as e:
-            node_rank = 0
-
-        # recover original exp before went into process
-        # init in write mode only on proc 0
-        self.experiment.debug = self.proc_rank > 0
-        self.experiment = self.experiment.get_non_ddp_exp()
-
-        # show progbar only on prog_rank 0
-        self.prog_bar = self.prog_bar and node_rank == 0 and gpu_nb == 0
-
-        # determine which process we are and world size
-        self.proc_rank = node_rank * len(self.data_parallel_device_ids) + gpu_nb
-        self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
-
-        # set up server using proc 0's ip address
-        # try to init for 20 times at max in case ports are taken
-        # where to store ip_table
-        self.__init_tcp_connection()
-
-        # CHOOSE OPTIMIZER
-        # filter out the weights that were done on gpu so we can load on good old cpus
-        self.optimizers = model.configure_optimizers()
-
-        # MODEL
-        # copy model to each gpu
-        torch.cuda.set_device(gpu_nb)
-        model.cuda(gpu_nb)
-
-        # AMP
-        # run through amp wrapper before going to distributed DP
-        if self.use_amp:
-            # An example
-            model, optimizers = amp.initialize(
-                model, self.optimizers, opt_level=self.amp_level,
-            )
-            self.optimizers = optimizers
-
-        model = LightningDistributedDataParallel(model, device_ids=[gpu_nb])
-
-        # continue training routine
-        self.__run_pretrain_routine(model)
-
-    def __init_tcp_connection(self):
-        """
-        Connect all procs in the world using the env:// init
-        Use the first node as the root address
-        :param port:
-        :param tries:
-        :return:
-        """
-        try:
-            port = os.environ['MASTER_PORT']
-        except Exception as e:
-            port = 12910
-            os.environ['MASTER_PORT'] = f'{port}'
-
-        try:
-            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
-        except Exception as e:
-            root_node = '127.0.0.2'
-
-        os.environ['MASTER_ADDR'] = root_node
-
-        sleep(self.proc_rank*0.5)
-        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
-
-    def __run_pretrain_routine(self, model):
-        """
-        Sanity check a few things before starting actual training
-        :param model:
-        :return:
-        """
-        ref_model = model
-        if self.data_parallel:
-            ref_model = model.module
-
-        ref_model.trainer = self
-
-        # set local properties on the model
-        ref_model.on_gpu = self.on_gpu
-
-        # transfer data loaders from model
-        self.__get_dataloaders(ref_model)
-
-        # init training constants
-        self.__layout_bookeeping()
-
-        # add lr schedulers
-        if self.lr_scheduler_milestones is not None:
-            for optimizer in self.optimizers:
-                scheduler = MultiStepLR(optimizer, self.lr_scheduler_milestones)
-                self.lr_schedulers.append(scheduler)
-
-        # print model summary
-        if self.proc_rank == 0 and self.print_weights_summary:
-            ref_model.summarize()
-
-        # give model convenience properties
-        ref_model.trainer = self
-        ref_model.experiment = self.experiment
-
-        # run tiny validation to make sure program won't crash during val
-        _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
-
-        # save exp to get started
-        if self.proc_rank == 0:
-            self.experiment.save()
-
-        # enable cluster checkpointing
-        if self.cluster is not None:
-            self.enable_auto_hpc_walltime_manager()
-
-        # ---------------------------
-        # CORE TRAINING LOOP
-        # ---------------------------
-        self.model = model
-        self.__train()
-
-    def __train(self):
-        # run all epochs
-        for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
-            # update the lr scheduler
-            for lr_scheduler in self.lr_schedulers:
-                lr_scheduler.step()
-
-            model = self.__get_model()
-            model.current_epoch = epoch_nb
-
-            # hook
-            if self.__is_function_implemented('on_epoch_start'):
-                model = self.__get_model()
-                model.on_epoch_start()
-
-            self.current_epoch = epoch_nb
-            self.total_batches = self.nb_tng_batches + self.nb_val_batches
-            self.batch_loss_value = 0  # accumulated grads
-
-            # init progbar when requested
-            if self.progress_bar:
-                self.prog_bar = tqdm.tqdm(range(self.total_batches), position=self.process_position)
-
-            for batch_nb, data_batch in enumerate(self.tng_dataloader):
-                self.batch_nb = batch_nb
-                self.global_step += 1
-
-                model = self.__get_model()
-                model.global_step = self.global_step
-
-                # stop when the flag is changed or we've gone past the amount requested in the batches
-                self.total_batch_nb += 1
-                met_batch_limit = batch_nb > self.nb_tng_batches
-                if met_batch_limit:
-                    break
-
-                # ---------------
-                # RUN TRAIN STEP
-                # ---------------
-                batch_result = self.__run_tng_batch(data_batch, batch_nb)
-                early_stop_epoch = batch_result == -1
-
-                # ---------------
-                # RUN VAL STEP
-                # ---------------
-                is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
-                if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
-                    self.__run_validation()
-
-                # when batch should be saved
-                if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
-                    if self.proc_rank == 0:
-                        self.experiment.save()
-
-                # when metrics should be logged
-                if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
-                    # count items in memory
-                    # nb_params, nb_tensors = count_mem_items()
-
-                    model = self.__get_model()
-                    metrics = model.update_tng_log_metrics(self.__tng_tqdm_dic)
-
-                    # add gpu memory
-                    if self.on_gpu:
-                        mem_map = get_gpu_memory_map()
-                        metrics.update(mem_map)
-
-                    # add norms
-                    if self.track_grad_norm > 0:
-                        model = self.__get_model()
-                        grad_norm_dic = model.grad_norm(self.track_grad_norm)
-                        metrics.update(grad_norm_dic)
-
-                    if self.__is_function_implemented('on_tng_metrics'):
-                        model.on_tng_metrics(metrics)
-
-                    # log metrics
-                    scalar_metrics = self.__metrics_to_scalars(metrics, blacklist=self.__log_vals_blacklist())
-                    if self.proc_rank == 0:
-                        self.experiment.log(scalar_metrics, global_step=self.global_step)
-                        self.experiment.save()
-
-                # hook
-                if self.__is_function_implemented('on_batch_end'):
-                    model = self.__get_model()
-                    model.on_batch_end()
-
-                # end epoch early
-                if early_stop_epoch:
-                    break
-
-            # hook
-            if self.__is_function_implemented('on_epoch_end'):
-                model = self.__get_model()
-                model.on_epoch_end()
-
-            # early stopping
-            met_min_epochs = epoch_nb > self.min_nb_epochs
-            if self.enable_early_stop and met_min_epochs:
-                should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, logs=self.__tng_tqdm_dic)
-
-                # stop training
-                stop = should_stop and met_min_epochs
-                if stop:
-                    return
-
-    def __metrics_to_scalars(self, metrics, blacklist=[]):
-        new_metrics = {}
-        for k, v in metrics.items():
-            if type(v) is torch.Tensor:
-                v = v.item()
-
-            if type(v) is dict:
-                v = self.__metrics_to_scalars(v)
-
-            if k not in blacklist:
-                new_metrics[k] = float(v)
-
-        return new_metrics
-
-    def __log_vals_blacklist(self):
-        """avoid logging some vals lightning uses to maintain state"""
-        blacklist = {'batch_nb', 'v_nb', 'gpu'}
-        return blacklist
-
-    def __run_tng_batch(self, data_batch, batch_nb):
-        if data_batch is None:
-            return 0
-
-        # hook
-        if self.__is_function_implemented('on_batch_start'):
-            model_ref = self.__get_model()
-            response = model_ref.on_batch_start(data_batch)
-
-            if response == -1:
-                return -1
-
-        if self.progress_bar:
-            self.prog_bar.update(1)
-
-        # forward pass
-        # return a scalar value and a dic with tqdm metrics
-        if self.use_ddp:
-            output = self.model(data_batch, batch_nb)
-        elif self.use_dp:
-            output = self.model(data_batch, batch_nb)
-            output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
-        else:
-            output = self.model.training_step(data_batch, batch_nb)
-
-        try:
-            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
-        except Exception as e:
-            model_specific_tqdm_metrics_dic = {}
-
-        # if output dict doesn't have the keyword loss
-        # then assume the output=loss if scalar
-        try:
-            loss = output['loss']
-        except Exception as e:
-            if type(output) is torch.Tensor:
-                loss = output
-
-        self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
-
-        # backward pass
-        if self.use_amp:
-            # scale loss when using amp
-            for optimizer in self.optimizers:
-                with amp.scale_loss(loss, optimizer) as scaled_loss:
-                    scaled_loss.backward()
-        else:
-            loss.backward()
-
-        if self.print_nan_grads:
-            model = self.__get_model()
-            for param in model.parameters():
-                print(param.grad.float().sum())
-
-        # avoid memory leaks
-        self.batch_loss_value += loss.item()
-
-        # gradient update with accumulated gradients
-        if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
-
-            # clip gradients
-            if self.gradient_clip > 0:
-                model = self.__get_model()
-                torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
-
-            # update gradients across all optimizers
-            for optimizer in self.optimizers:
-                optimizer.step()
-
-                # clear gradients
-                optimizer.zero_grad()
-
-            # queuing loss across batches blows it up proportionally... divide out the number accumulated
-            self.batch_loss_value = self.batch_loss_value / self.accumulate_grad_batches
-
-            # track loss
-            self.running_loss.append(self.batch_loss_value)
-            self.batch_loss_value = 0
-            self.avg_loss = np.mean(self.running_loss[-100:])
-
-            # update progbar
-            if self.progress_bar:
-                # add model specific metrics
-                tqdm_metrics = self.__tng_tqdm_dic
-                self.prog_bar.set_postfix(**tqdm_metrics)
-
-        # activate batch end hook
-        if self.__is_function_implemented('on_batch_end'):
-            model = self.__get_model()
-            model.on_batch_end()
-
-        return 0
-
-    def __run_validation(self):
-        # decide if can check epochs
-        can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
-        if self.fast_dev_run:
-            print('skipping to check performance bc of --fast_dev_run')
-        elif not can_check_epoch:
-            return
-
-        try:
-            # hook
-            if self.__is_function_implemented('on_pre_performance_check'):
-                model = self.__get_model()
-                model.on_pre_performance_check()
-
-            # use full val set on end of epoch
-            # use a small portion otherwise
-            max_batches = None if not self.fast_dev_run else 1
-            model_specific_tqdm_metrics_dic = self.validate(
-                self.model,
-                self.val_dataloader,
-                max_batches
-            )
-            self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
-
-            # hook
-            if self.__is_function_implemented('on_post_performance_check'):
-                model = self.__get_model()
-                model.on_post_performance_check()
-
-        except Exception as e:
-            print(e)
-            print(traceback.print_exc())
-
-        if self.progress_bar:
-            # add model specific metrics
-            tqdm_metrics = self.__tng_tqdm_dic
-            self.prog_bar.set_postfix(**tqdm_metrics)
-
-        # model checkpointing
-        if self.proc_rank == 0 and self.checkpoint_callback:
-            print('save callback...')
-            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -6,10 +6,9 @@ from itertools import chain
 import threading
 import torch
 from torch.cuda._utils import _get_device_index
-import pdb
 
 
-def _find_tensors(obj):
+def _find_tensors(obj):  # pragma: no cover
     r"""
     Recursively find all tensors contained in the specified object.
     """
@@ -22,8 +21,7 @@ def _find_tensors(obj):
     return []
 
 
-
-def get_a_var(obj):
+def get_a_var(obj):  # pragma: no cover
     if isinstance(obj, torch.Tensor):
         return obj
 
@@ -58,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])
 
@@ -65,7 +65,6 @@ class LightningDataParallel(DataParallel):
         outputs = self.parallel_apply(replicas, inputs, kwargs)
         return self.gather(outputs, self.output_device)
 
-
     def parallel_apply(self, replicas, inputs, kwargs):
         return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
 
@@ -78,7 +77,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
     def parallel_apply(self, replicas, inputs, kwargs):
         return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
 
-    def forward(self, *inputs, **kwargs):
+    def forward(self, *inputs, **kwargs):  # pragma: no cover
         self._sync_params()
         if self.device_ids:
             inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
@@ -88,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():
@@ -113,7 +114,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
         return output
 
 
-def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
+def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):  # pragma: no cover
     r"""Applies each `module` in :attr:`modules` in parallel on arguments
     contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
     on each of :attr:`devices`.
@@ -156,6 +157,10 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
                 # 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)
                 # ---------------
@@ -166,6 +171,14 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
             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

@@ -0,0 +1,33 @@
+import traceback
+
+
+def data_loader(fn):
+    """
+    Decorator to make any fx with this use the lazy property
+    :param fn:
+    :return:
+    """
+
+    attr_name = '_lazy_' + fn.__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 _get_data_loader

pytorch_lightning/root_module/grads.py

@@ -1,10 +1,9 @@
-import numpy as np
-from torch import nn
-
 """
 Module to describe gradients
 """
 
+from torch import nn
+
 
 class GradInformation(nn.Module):
 
@@ -18,23 +17,13 @@ class GradInformation(nn.Module):
                     total_norm += param_norm ** norm_type
                     norm = param_norm ** (1 / norm_type)
 
-                    results['grad_{}_norm_{}'.format(norm_type, i)] = round(norm.data.cpu().numpy().flatten()[0], 3)
-                except Exception as e:
+                    grad = round(norm.data.cpu().numpy().flatten()[0], 3)
+                    results['grad_{}_norm_{}'.format(norm_type, i)] = grad
+                except Exception:
                     # this param had no grad
                     pass
 
         total_norm = total_norm ** (1. / norm_type)
-        results['grad_{}_norm_total'.format(norm_type)] = round(total_norm.data.cpu().numpy().flatten()[0], 3)
+        grad = round(total_norm.data.cpu().numpy().flatten()[0], 3)
+        results['grad_{}_norm_total'.format(norm_type)] = grad
         return results
-
-
-    def describe_grads(self):
-        for p in self.parameters():
-            g = p.grad.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))
-
-
-    def describe_params(self):
-        for p in self.parameters():
-            g = p.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))

pytorch_lightning/root_module/hooks.py

@@ -1,7 +1,16 @@
 import torch
 
+
 class ModelHooks(torch.nn.Module):
-    def on_batch_start(self, data_batch):
+
+    def on_sanity_check_start(self):
+        """
+        Called before starting evaluate
+        :return:
+        """
+        pass
+
+    def on_batch_start(self, batch):
         pass
 
     def on_batch_end(self):
@@ -19,6 +28,23 @@ class ModelHooks(torch.nn.Module):
     def on_post_performance_check(self):
         pass
 
-    def on_tng_metrics(self, metrics):
+    def on_before_zero_grad(self, optimizer):
+        """
+        Called after optimizer.step() and before optimizer.zero_grad()
+
+        for optimizer in optimizers:
+            optimizer.step()
+            model.on_before_zero_grad(optimizer) # < ---- called here
+            optimizer.zero_grad
+
+        :param optimizer:
+        :return:
+        """
         pass
 
+    def on_after_backward(self):
+        """
+        Called after loss.backward() and before optimizers do anything
+        :return:
+        """
+        pass

pytorch_lightning/root_module/memory.py

@@ -1,14 +1,14 @@
-import torch
-import gc
-import subprocess
-import numpy as np
-import pandas as pd
-
-
 '''
 Generates a summary of a model's layers and dimensionality
 '''
 
+import gc
+
+import torch
+import subprocess
+import numpy as np
+import pandas as pd
+
 
 class ModelSummary(object):
 
@@ -33,33 +33,42 @@ class ModelSummary(object):
         mods = list(self.model.modules())
         in_sizes = []
         out_sizes = []
-        input_ = self.example_input_array
-        for i in range(1, len(mods)):
-            m = mods[i]
-            if type(input_) is list or type(input_) is tuple:
-                out = m(*input_)
-            else:
-                out = m(input_)
+        input_ = self.model.example_input_array
 
-            if type(input_) is tuple or type(input_) is list:
-                in_size = []
-                for x in input_:
-                    if type(x) is list:
-                        in_size.append(len(x))
-                    else:
-                        in_size.append(x.size())
-            else:
-                in_size = np.array(input_.size())
+        if self.model.on_gpu:
+            input_ = input_.cuda(0)
 
-            in_sizes.append(in_size)
+        if self.model.trainer.use_amp:
+            input_ = input_.half()
 
-            if type(out) is tuple or type(out) is list:
-                out_size = np.asarray([x.size() for x in out])
-            else:
-                out_size = np.array(out.size())
+        with torch.no_grad():
 
-            out_sizes.append(out_size)
-            input_ = out
+            for i in range(1, len(mods)):
+                m = mods[i]
+                if type(input_) is list or type(input_) is tuple:  # pragma: no cover
+                    out = m(*input_)
+                else:
+                    out = m(input_)
+
+                if type(input_) is tuple or type(input_) is list:  # pragma: no cover
+                    in_size = []
+                    for x in input_:
+                        if type(x) is list:
+                            in_size.append(len(x))
+                        else:
+                            in_size.append(x.size())
+                else:
+                    in_size = np.array(input_.size())
+
+                in_sizes.append(in_size)
+
+                if type(out) is tuple or type(out) is list:  # pragma: no cover
+                    out_size = np.asarray([x.size() for x in out])
+                else:
+                    out_size = np.array(out.size())
+
+                out_sizes.append(out_size)
+                input_ = out
 
         self.in_sizes = in_sizes
         self.out_sizes = out_sizes
@@ -85,7 +94,7 @@ class ModelSummary(object):
         mods = list(self.model.modules())
         sizes = []
 
-        for i in range(1,len(mods)):
+        for i in range(1, len(mods)):
             m = mods[i]
             p = list(m.parameters())
             modsz = []
@@ -114,13 +123,22 @@ class ModelSummary(object):
         Layer Name, Layer Type, Input Size, Output Size, Number of Parameters
         '''
 
-        df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) )
-        df.columns = ['Name', 'Type', 'Params']
+        cols = ['Name', 'Type', 'Params']
+        if self.model.example_input_array is not None:
+            cols.extend(['In_sizes', 'Out_sizes'])
+
+        df = pd.DataFrame(np.zeros((len(self.layer_names), len(cols))))
+        df.columns = cols
 
         df['Name'] = self.layer_names
         df['Type'] = self.layer_types
         df['Params'] = self.param_nums
 
+        if self.model.example_input_array is not None:
+
+            df['In_sizes'] = self.in_sizes
+            df['Out_sizes'] = self.out_sizes
+
         self.summary = df
         return
 
@@ -128,19 +146,22 @@ class ModelSummary(object):
         self.get_layer_names()
         self.get_parameter_sizes()
         self.get_parameter_nums()
+
+        if self.model.example_input_array is not None:
+            self.get_variable_sizes()
         self.make_summary()
 
 
-def print_mem_stack():
+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())
-        except Exception as e:
+        except Exception:
             pass
 
 
-def count_mem_items():
+def count_mem_items():  # pragma: no cover
     nb_params = 0
     nb_tensors = 0
     for obj in gc.get_objects():
@@ -151,12 +172,39 @@ def count_mem_items():
                     nb_params += 1
                 else:
                     nb_tensors += 1
-        except Exception as e:
+        except Exception:
             pass
 
     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_mem = 1000000
+        min_k = None
+        max_mem = 0
+        max_k = None
+        for k, v in memory_map:
+            if v > max_mem:
+                max_mem = v
+                max_k = k
+            if v < min_mem:
+                min_mem = v
+                min_k = k
+
+        memory_map = {min_k: min_mem, max_k: max_mem}
+
+    return memory_map
+
+
 def get_gpu_memory_map():
     """Get the current gpu usage.
 

pytorch_lightning/root_module/model_saving.py

@@ -1,206 +1,34 @@
-import torch
-import os
-import re
-import pdb
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
-
 class ModelIO(object):
 
-    def load_model_specific(self, checkpoint):
+    def on_load_checkpoint(self, checkpoint):
         """
         Do something with the checkpoint
+        Gives model a chance to load something before state_dict is restored
         :param checkpoint:
         :return:
         """
-        raise NotImplementedError
+        pass
 
-    def get_save_dict(self):
+    def on_save_checkpoint(self, checkpoint):
         """
-        Return specific things for the model
-        :return:
+        Give the model a chance to add something to the checkpoint.
+        state_dict is already there
         """
-        raise NotImplementedError
+        pass
 
     # -------------------------
     # OPTIONAL HOOKS
     # -------------------------
-    def on_hpc_save(self):
+    def on_hpc_save(self, checkpoint):
         """
         Hook to do whatever you need right before Slurm manager saves the model
         :return:
         """
         pass
 
-    def on_hpc_load(self):
+    def on_hpc_load(self, checkpoint):
         """
         Hook to do whatever you need right before Slurm manager loads the model
         :return:
         """
         pass
-
-
-class TrainerIO(object):
-
-    # --------------------
-    # MODEL SAVE CHECKPOINT
-    # --------------------
-    def save_checkpoint(self, filepath):
-        checkpoint = self.dump_checkpoint()
-
-        # do the actual save
-        torch.save(checkpoint, filepath)
-
-    def dump_checkpoint(self):
-        checkpoint = {
-            'epoch': self.current_epoch,
-            'checkpoint_callback_best': self.checkpoint_callback.best,
-            'early_stop_callback_wait': self.early_stop_callback.wait,
-            'early_stop_callback_patience': self.early_stop_callback.patience,
-            'global_step': self.global_step
-        }
-
-        optimizer_states = []
-        for i, optimizer in enumerate(self.optimizers):
-            optimizer_states.append(optimizer.state_dict())
-
-        checkpoint['optimizer_states'] = optimizer_states
-
-        # request what to save from the model
-        is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
-        model = self.model.module if is_dp_module else self.model
-        checkpoint_dict = model.get_save_dict()
-
-        # merge trainer and model saving items
-        checkpoint.update(checkpoint_dict)
-        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
-        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:
-        """
-        self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
-        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)
-
-    # ----------------------------------
-    # 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
-        self.on_hpc_save()
-
-        # request what to save from the model
-        checkpoint_dict = self.dump_checkpoint()
-
-        # do the actual save
-        torch.save(checkpoint_dict, 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
-        self.restore_training_state(checkpoint)
-
-        # load model state
-        model = self.model.module if type(self.model) is LightningDataParallel else self.model
-        model.load_model_specific(checkpoint)
-
-        # call model hook
-        self.on_hpc_load()
-
-    def max_ckpt_in_folder(self, path):
-        files = os.listdir(path)
-        files = [x for x in files if 'ckpt_' in x]
-        if len(files) == 0:
-            return 0
-
-        ckpt_vs = []
-        for name in files:
-            name = name.split('ckpt_')[-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/optimization.py

@@ -1,22 +0,0 @@
-from torch import nn
-from torch import optim
-
-
-class OptimizerConfig(nn.Module):
-
-    def choose_optimizer(self, optimizer, params, optimizer_params, opt_name_key):
-        if optimizer == 'adam':
-            optimizer = optim.Adam(params, **optimizer_params)
-        if optimizer == 'sparse_adam':
-            optimizer = optim.SparseAdam(params, **optimizer_params)
-        if optimizer == 'sgd':
-            optimizer = optim.SGD(params, **optimizer_params)
-        if optimizer == 'adadelta':
-            optimizer = optim.Adadelta(params, **optimizer_params)
-
-        # transfer opt state if loaded
-        if opt_name_key in self.loaded_optimizer_states_dict:
-            state = self.loaded_optimizer_states_dict[opt_name_key]
-            optimizer.load_state_dict(state)
-
-        return optimizer

pytorch_lightning/root_module/root_module.py

@@ -1,19 +1,17 @@
-import os
 import torch
-import math
 
 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.optimization import OptimizerConfig
+from pytorch_lightning.trainer.trainer_io import load_hparams_from_tags_csv
+from pytorch_lightning.root_module.model_saving import ModelIO
 from pytorch_lightning.root_module.hooks import ModelHooks
+from pytorch_lightning.root_module.decorators import data_loader
 
 
-class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
+class LightningModule(GradInformation, ModelIO, ModelHooks):
 
-    def __init__(self, hparams):
-        super(LightningModule, self).__init__()
-        self.hparams = hparams
+    def __init__(self, *args, **kwargs):
+        super(LightningModule, self).__init__(*args, **kwargs)
 
         self.dtype = torch.FloatTensor
         self.exp_save_path = None
@@ -21,15 +19,15 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, 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
-
-        # computed vars for the dataloaders
-        self._tng_dataloader = None
-        self._val_dataloader = None
-        self._test_dataloader = None
+        self.use_dp = False
+        self.use_ddp = False
+        self.use_ddp2 = False
+        self.use_amp = False
 
     def forward(self, *args, **kwargs):
         """
@@ -40,59 +38,131 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, 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 loss, dict with metrics for tqdm
+        :param called with batch, batch_nb
+        additional: optimizer_i if multiple optimizers used
         :return:
         """
         raise NotImplementedError
 
+    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_optimizers(self):
         """
-        Return array of optimizers
+        Return a list of optimizers and a list of schedulers (could be empty)
         :return:
         """
         raise NotImplementedError
 
-    def update_tng_log_metrics(self, logs):
+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, second_order_closure=None):
         """
-        Chance to update metrics to be logged for training step.
-        For example, add music, images, etc... to log
-        :param logs:
+        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:
         """
-        return logs
+        if isinstance(optimizer, torch.optim.LBFGS):
+            optimizer.step(second_order_closure)
+        else:
+            optimizer.step()
 
-    def loss(self, *args, **kwargs):
+        # clear gradients
+        optimizer.zero_grad()
+
+    @data_loader
+    def train_dataloader(self):
         """
-        Expand model_out into your components
-        :param model_out:
+        Implement a PyTorch DataLoader
         :return:
         """
         raise NotImplementedError
 
+    @data_loader
+    def test_dataloader(self):
+        """
+        Implement a PyTorch DataLoader
+        :return:
+        """
+        return None
+
+    @data_loader
+    def val_dataloader(self):
+        """
+        Implement a PyTorch DataLoader
+        :return:
+        """
+        return None
+
+    @classmethod
+    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 map_location: dic for mapping storage {'cuda:1':'cuda:0'}
+        :return:
+        """
+        hparams = load_hparams_from_tags_csv(tags_csv)
+        hparams.__setattr__('on_gpu', False)
+
+        # 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)
+        model.load_state_dict(checkpoint['state_dict'])
+
+        # give model a chance to load something
+        model.on_load_checkpoint(checkpoint)
+
+        return model
+
     def summarize(self):
         model_summary = ModelSummary(self)
         print(model_summary)
 
-
     def freeze(self):
         for param in self.parameters():
             param.requires_grad = False
@@ -100,65 +170,3 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
     def unfreeze(self):
         for param in self.parameters():
             param.requires_grad = True
-
-    @property
-    def tng_dataloader(self):
-        """
-        Implement a function to load an h5py of this data
-        :return:
-        """
-        raise NotImplementedError
-
-    @property
-    def test_dataloader(self):
-        """
-        Implement a function to load an h5py of this data
-        :return:
-        """
-        raise NotImplementedError
-
-    @property
-    def val_dataloader(self):
-        """
-        Implement a function to load an h5py of this data
-        :return:
-        """
-        raise NotImplementedError
-
-    @staticmethod
-    def get_process_position(gpus):
-        try:
-            current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
-            gpu_ids = gpus.split(',')
-            process_position = gpu_ids.index(current_gpu)
-            return process_position, current_gpu
-        except Exception as e:
-            return 0, 0
-
-    @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
-        """
-        Primary way of loading model from csv weights path
-        :param weights_path:
-        :param tags_csv:
-        :param on_gpu:
-        :param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
-        :return:
-        """
-        hparams = load_hparams_from_tags_csv(tags_csv)
-        hparams.__setattr__('on_gpu', on_gpu)
-
-        if on_gpu:
-            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)
-
-        model = cls(hparams)
-
-        # allow model to load
-        model.load_model_specific(checkpoint)
-        model.load_state_dict(checkpoint['state_dict'], strict=False)
-        return model

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

@@ -0,0 +1,27 @@
+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(LightningValidationMixin, LightningTestMixin, LightningTestModelBase):
+    """
+    Most common test case. Validation and test dataloaders
+    """
+
+    def on_training_metrics(self, logs):
+        logs['some_tensor_to_test'] = torch.rand(1)

pytorch_lightning/testing/lm_test_module_base.py

@@ -0,0 +1,198 @@
+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
+
+
+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}
+            })
+
+            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,391 @@
+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
+
+
+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:
+                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
+
+        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}
+        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/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/trainer_io.py

@@ -0,0 +1,325 @@
+import os
+import re
+import signal
+import pdb
+from subprocess import call
+
+import torch
+
+from pytorch_lightning.pt_overrides.override_data_parallel import (
+    LightningDistributedDataParallel, LightningDataParallel)
+
+
+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
+
+    # --------------------
+    # CHECK-POINTING
+    # --------------------
+    def restore_weights(self, model):
+        """
+        To restore weights we have two cases.
+        First, if we use the same experiment version, then restore the latest ckpt.
+        AFTER that, if we find weights from hpc checkpoint, then restore that.
+        :param model:
+        :return:
+        """
+        # restore weights if same exp version
+        self.restore_state_if_checkpoint_exists(model)
+
+        # if script called from hpc resubmit, load weights
+        self.restore_hpc_weights_if_needed(model)
+
+    def restore_state_if_checkpoint_exists(self, model):
+        # do nothing if there's not dir or callback
+        no_ckpt_callback = self.checkpoint_callback is None
+        if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath):
+            return
+
+        # 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)
+            print(f'model and trainer restored from checkpoint: {last_ckpt_path}')
+
+    # --------------------
+    # 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:
+            print('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
+            print('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
+            print('\nrequeing job {}...'.format(job_id))
+            result = call(cmd, shell=True)
+
+            # print result text
+            if result == 0:
+                print('requeued exp ', job_id)
+            else:
+                print('requeue failed...')
+
+            # close experiment to avoid issues
+            self.logger.close()
+
+    def term_handler(self, signum, frame):
+        # save
+        print("bypassing sigterm")
+
+    # --------------------
+    # 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:
+        # 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:
+            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 restore_hpc_weights_if_needed(self, model):
+        """
+        If there is a set of hpc weights, use as signal to restore model
+        :param model:
+        :return:
+        """
+        # 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)
+
+    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)
+
+            # 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
+        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)
+
+        print(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_main.py

@@ -1,214 +0,0 @@
-import os
-import sys
-
-import torch
-import numpy as np
-from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
-from time import sleep
-
-from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
-
-# ---------------------
-# DEFINE MODEL HERE
-# ---------------------
-from pytorch_lightning.models.sample_model_template.model_template import ExampleModel1
-# ---------------------
-
-AVAILABLE_MODELS = {
-    'model_1': ExampleModel1
-}
-
-
-"""
-Allows training by using command line arguments
-
-Run by: 
-# TYPE YOUR RUN COMMAND HERE
-"""
-
-
-def main_local(hparams):
-    main(hparams, None, None)
-
-
-def main(hparams, cluster, results_dict):
-    """
-    Main training routine specific for this project
-    :param hparams:
-    :return:
-    """
-    on_gpu = torch.cuda.is_available()
-    if hparams.disable_cuda:
-        on_gpu = False
-
-    device = 'cuda' if on_gpu else 'cpu'
-    hparams.__setattr__('device', device)
-    hparams.__setattr__('on_gpu', on_gpu)
-    hparams.__setattr__('nb_gpus', torch.cuda.device_count())
-    hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
-
-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
-    # init experiment
-    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
-    print('loading model...')
-    model = TRAINING_MODEL(hparams)
-    print('model built')
-
-    # callbacks
-    early_stop = EarlyStopping(
-        monitor=hparams.early_stop_metric,
-        patience=hparams.early_stop_patience,
-        verbose=True,
-        mode=hparams.early_stop_mode
-    )
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-    checkpoint = ModelCheckpoint(
-        filepath=model_save_path,
-        save_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor=hparams.model_save_monitor_value,
-        mode=hparams.model_save_monitor_mode
-    )
-
-    # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        on_gpu=on_gpu,
-        cluster=cluster,
-        progress_bar=hparams.enable_tqdm,
-        overfit_pct=hparams.overfit,
-        track_grad_norm=hparams.track_grad_norm,
-        fast_dev_run=hparams.fast_dev_run,
-        check_val_every_n_epoch=hparams.check_val_every_n_epoch,
-        accumulate_grad_batches=hparams.accumulate_grad_batches,
-        process_position=process_position,
-        current_gpu_name=current_gpu,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-        enable_early_stop=hparams.enable_early_stop,
-        max_nb_epochs=hparams.max_nb_epochs,
-        min_nb_epochs=hparams.min_nb_epochs,
-        train_percent_check=hparams.train_percent_check,
-        val_percent_check=hparams.val_percent_check,
-        test_percent_check=hparams.test_percent_check,
-        val_check_interval=hparams.val_check_interval,
-        log_save_interval=hparams.log_save_interval,
-        add_log_row_interval=hparams.add_log_row_interval,
-        lr_scheduler_milestones=hparams.lr_scheduler_milestones
-    )
-
-    # train model
-    trainer.fit(model)
-
-
-def get_default_parser(strategy, root_dir):
-
-    possible_model_names = list(AVAILABLE_MODELS.keys())
-    parser = HyperOptArgumentParser(strategy=strategy, add_help=False)
-    add_default_args(parser, root_dir, possible_model_names, SEED)
-    return parser
-
-
-def get_model_name(args):
-    for i, arg in enumerate(args):
-        if 'model_name' in arg:
-            return args[i+1]
-
-
-def optimize_on_cluster(hyperparams):
-    # enable cluster training
-    cluster = SlurmCluster(
-        hyperparam_optimizer=hyperparams,
-        log_path=hyperparams.tt_save_path,
-        test_tube_exp_name=hyperparams.tt_name
-    )
-
-    # email for cluster coms
-    cluster.notify_job_status(email='add_email_here', on_done=True, on_fail=True)
-
-    # configure cluster
-    cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus
-    cluster.job_time = '48:00:00'
-    cluster.gpu_type = '1080ti'
-    cluster.memory_mb_per_node = 48000
-
-    # any modules for code to run in env
-    cluster.add_command('source activate pytorch_lightning')
-
-    # name of exp
-    job_display_name = hyperparams.tt_name.split('_')[0]
-    job_display_name = job_display_name[0:3]
-
-    # run hopt
-    print('submitting jobs...')
-    cluster.optimize_parallel_cluster_gpu(
-        main,
-        nb_trials=hyperparams.nb_hopt_trials,
-        job_name=job_display_name
-    )
-
-
-if __name__ == '__main__':
-
-    model_name = get_model_name(sys.argv)
-
-    # use default args
-    root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
-    parent_parser = get_default_parser(strategy='random_search', root_dir=root_dir)
-
-    # allow model to overwrite or extend args
-    TRAINING_MODEL = AVAILABLE_MODELS[model_name]
-    parser = TRAINING_MODEL.add_model_specific_args(parent_parser)
-    parser.json_config('-c', '--config', default=root_dir + '/run_configs/local.json')
-    hyperparams = parser.parse_args()
-
-    # format GPU layout
-    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-    gpu_ids = hyperparams.gpus.split(';')
-
-    # RUN TRAINING
-    if hyperparams.on_cluster:
-        print('RUNNING ON SLURM CLUSTER')
-        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
-        optimize_on_cluster(hyperparams)
-
-    elif hyperparams.single_run_gpu:
-        print(f'RUNNING 1 TRIAL ON GPU. gpu: {gpu_ids[0]}')
-        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[0]
-        main(hyperparams, None, None)
-
-    elif hyperparams.local or hyperparams.single_run:
-        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
-        print('RUNNING LOCALLY')
-        main(hyperparams, None, None)
-
-    else:
-        print(f'RUNNING MULTI GPU. GPU ids: {gpu_ids}')
-        hyperparams.optimize_parallel_gpu(
-            main_local,
-            gpu_ids=gpu_ids,
-            nb_trials=hyperparams.nb_hopt_trials,
-            nb_workers=len(gpu_ids)
-        )

pytorch_lightning/utilities/arg_parse.py

@@ -1,31 +1,48 @@
-import pdb
+"""
+List of default args which mught be useful for all the available flags
+Might need to update with the new flags
+"""
+
+import os
+
 
 def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None):
 
-    # tng, test, val check intervals
-    parser.add_argument('--eval_test_set', dest='eval_test_set', action='store_true', help='true = run test set also')
-    parser.add_argument('--check_val_every_n_epoch', default=1, type=int, help='check val every n epochs')
+    # 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,
+                        help='check val every n epochs')
     parser.opt_list('--accumulate_grad_batches', default=1, type=int, tunable=False,
-                    help='accumulates gradients k times before applying update. Simulates huge batch size')
+                    help='accumulates gradients k times before applying update.'
+                         ' Simulates huge batch size')
     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')
-    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, help='how much of test set to check')
+    parser.add_argument('--train_percent_check', default=1.0, type=float,
+                        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,
+                        help='how much of test set to check')
 
-    parser.add_argument('--val_check_interval', default=0.95, type=float, 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, help='add log every k batches')
+    parser.add_argument('--val_check_interval', default=0.95, type=float,
+                        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('--row_log_interval', default=100, type=int,
+                        help='add log every k batches')
 
     # early stopping
     parser.add_argument('--disable_early_stop', dest='enable_early_stop', action='store_false')
     parser.add_argument('--early_stop_metric', default='val_acc', type=str)
     parser.add_argument('--early_stop_mode', default='min', type=str)
-    parser.add_argument('--early_stop_patience', default=3, type=int, help='number of epochs until stop')
+    parser.add_argument('--early_stop_patience', default=3, type=int,
+                        help='number of epochs until stop')
 
     # gradient handling
-    parser.add_argument('--gradient_clip', default=-1, type=int)
-    parser.add_argument('--track_grad_norm', default=-1, type=int, help='if > 0, will track this grad norm')
+    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')
 
     # model saving
     parser.add_argument('--model_save_path', default=root_dir + '/model_weights')
@@ -41,7 +58,8 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     # test_tube settings
     parser.add_argument('-en', '--tt_name', default='pt_test')
     parser.add_argument('-td', '--tt_description', default='pytorch lightning test')
-    parser.add_argument('--tt_save_path', default=root_dir + '/test_tube_logs', help='logging dir')
+    parser.add_argument('--tt_save_path', default=os.path.join(root_dir, 'test_tube_logs'),
+                        help='logging dir')
     parser.add_argument('--enable_single_run', dest='single_run', action='store_true')
     parser.add_argument('--nb_hopt_trials', default=1, type=int)
     parser.add_argument('--log_stdout', dest='log_stdout', action='store_true')
@@ -52,25 +70,30 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     parser.add_argument('--default_tensor_type', default='torch.cuda.FloatTensor', type=str)
     parser.add_argument('--use_amp', dest='use_amp', action='store_true')
     parser.add_argument('--check_grad_nans', dest='check_grad_nans', action='store_true')
-    parser.add_argument('--amp_level', default='O2',type=str)
-
+    parser.add_argument('--amp_level', default='O2', type=str)
 
     # run on hpc
     parser.add_argument('--on_cluster', dest='on_cluster', action='store_true')
 
     # 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')
-    parser.add_argument('--enable_tqdm', dest='enable_tqdm', default=False, action='store_true', help='false removes the prog bar')
-    parser.add_argument('--overfit', default=-1, type=float, help='% of dataset to use with this option. float, or -1 for none')
+    parser.add_argument('--fast_dev_run', dest='fast_dev_run', default=False, action='store_true',
+                        help='runs validation after 1 training step')
+    parser.add_argument('--enable_tqdm', dest='enable_tqdm', default=False, action='store_true',
+                        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')
 
     # debug args
     if rand_seed is not None:
         parser.add_argument('--random_seed', default=rand_seed, type=int)
 
-    parser.add_argument('--interactive', dest='interactive', action='store_true', help='runs on gpu without cluster')
-    parser.add_argument('--debug', dest='debug', action='store_true', help='enables/disables test tube')
-    parser.add_argument('--local', dest='local', action='store_true', help='enables local tng')
+    parser.add_argument('--interactive', dest='interactive', action='store_true',
+                        help='runs on gpu without cluster')
+    parser.add_argument('--debug', dest='debug', action='store_true',
+                        help='enables/disables test tube')
+    parser.add_argument('--local', dest='local', action='store_true',
+                        help='enables local training')
 
     # optimizer
-    parser.add_argument('--lr_scheduler_milestones', default=None, type=str)
+    parser.add_argument('--lr_scheduler_milestones', default=None, type=str)

pytorch_lightning/utilities/debugging.py

@@ -0,0 +1,2 @@
+class MisconfigurationException(Exception):
+    pass

pytorch_lightning/utils/embeddings.py

@@ -1,104 +0,0 @@
-import torch
-import numpy as np
-from copy import deepcopy
-
-
-class PretrainedEmbedding(torch.nn.Embedding):
-
-    def __init__(self, embedding_path, embedding_dim, task_vocab, freeze=True, *args, **kwargs):
-        """
-        Loads a prebuilt pytorch embedding from any embedding formated file.
-        Padding=0 by default.
-
-        >>> emb = PretrainedEmbedding(embedding_path='glove.840B.300d.txt',embedding_dim=300, task_vocab={'hello': 1, 'world': 2})
-        >>> data = torch.Tensor([[0, 1], [0, 2]]).long()
-        >>> embedded = emb(data)
-        
-
-
-        :param embedding_path:
-        :param emb_dim:
-        :param task_vocab:
-        :param freeze:
-        :return:
-        """
-        # count the vocab
-        self.vocab_size = max(task_vocab.values()) + 1
-        super(PretrainedEmbedding, self).__init__(self.vocab_size, embedding_dim, padding_idx=0, *args, **kwargs)
-
-        # load pretrained embeddings
-        new_emb = self.__load_task_specific_embeddings(deepcopy(task_vocab), embedding_path, embedding_dim, freeze)
-
-        # transfer weights
-        self.weight = new_emb.weight
-
-        # apply freeze
-        should_freeze = not freeze
-        self.weight.requires_grad = should_freeze
-
-    def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
-        """
-        Iterates embedding file to only pull out task specific embeddings
-        :param vocab_words:
-        :param embedding_path:
-        :param emb_dim:
-        :param freeze:
-        :return:
-        """
-
-        # holds final embeddings for relevant words
-        embeddings = np.zeros(shape=(self.vocab_size, emb_dim))
-
-        # load embedding line by line and extract relevant embeddings
-        with open(embedding_path, encoding='utf-8') as f:
-            for line in f:
-                tokens = line.split(' ')
-                word = tokens[0]
-                embedding = tokens[1:]
-                embedding[-1] = embedding[-1][:-1]  # remove last new line
-
-                if word in vocab_words:
-                    vocab_word_i = vocab_words[word]
-
-                    # skip words that try to overwrite pad idx
-                    if vocab_word_i == 0:
-                        del vocab_words[word]
-                        continue
-
-                    emb_vals = np.asarray([float(x) for x in embedding])
-                    embeddings[vocab_word_i] = emb_vals
-
-                    # remove vocab word to early terminate
-                    del vocab_words[word]
-
-                # early break
-                if len(vocab_words) == 0:
-                    break
-
-        # add random vectors for the non-pretrained words
-        # these are vocab words NOT found in the pretrained embeddings
-        for w, i in vocab_words.items():
-            # skip words that try to overwrite pad idx
-            if i == 0:
-                continue
-
-            embedding = np.random.normal(size=emb_dim)
-            embeddings[i] = embedding
-
-        # turn into pt embedding
-        embeddings = torch.FloatTensor(embeddings)
-        embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=freeze)
-
-        return embeddings
-
-
-if __name__ == '__main__':
-    emb = PretrainedEmbedding(
-        embedding_path='/Users/waf/Developer',
-        embedding_dim=300,
-        task_vocab={'hello': 1, 'world': 2}
-    )
-
-    data = torch.Tensor([[0, 1], [0, 2]]).long()
-    embedded = emb(data)
-    print(embedded)

pytorch_lightning/utils/plotting.py

@@ -1,28 +0,0 @@
-import numpy as np
-np.seterr(divide='ignore', invalid='ignore')
-
-
-def plot_confusion_matrix(cm,
-                          save_path,
-                          normalize=False,
-                          title='Confusion matrix',
-                          ylabel='y',
-                          xlabel='x'):
-    """
-    This function prints and plots the confusion matrix.
-    Normalization can be applied by setting `normalize=True`.
-    """
-    from matplotlib import pyplot as plt
-    if normalize:
-        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
-        print("Normalized confusion matrix")
-    else:
-        print('Confusion matrix, without normalization')
-
-    fig = plt.figure()
-    plt.matshow(cm)
-    plt.title(title)
-    plt.colorbar()
-    plt.ylabel(ylabel)
-    plt.xlabel(xlabel)
-    plt.savefig(save_path)

requirements.txt

@@ -1,35 +1,7 @@
-
-atomicwrites==1.2.1
-attrs==18.2.0
-certifi==2018.11.29
-cffi==1.11.5
-imageio==2.4.1
-mkl-fft==1.0.6
-mkl-random==1.0.2
-more-itertools==5.0.0
-numpy==1.15.4
-olefile==0.46
-pandas==0.23.4
-Pillow==5.3.0
-pluggy==0.8.0
-py==1.7.0
-pycparser==2.19
-pytest==4.0.2
-python-dateutil==2.7.5
-pytz==2018.7
 scikit-learn==0.20.2
-scipy==1.2.0
-six==1.12.0
-sklearn==0.0
-tensorboard==1.14.0
-tensorboardX==1.7
-tensorflow==1.14.0
-test-tube==0.643
-torch==1.0.0
-torchvision==0.2.1
-tqdm==4.32.1
+tqdm==4.35.0
 twine==1.13.0
-urllib3==1.25.3
-webencodings==0.5.1
-Werkzeug==0.15.4
-wrapt==1.11.2
+numpy==1.16.4
+torch>=1.2.0
+torchvision>=0.3.0
+pandas

setup.cfg

@@ -16,6 +16,37 @@ markers =
 ignore = E731,W504
 max-line-length = 120
 
+[coverage:report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    if self.debug:
+    if settings.DEBUG
+    raise AssertionError
+    raise NotImplementedError
+    if 0:
+    if __name__ == .__main__.:
+    except Exception as e
+    print(e)
+    print(traceback.print_exc())
+    return *
+    raise Exception
+    raise *
+    except *
+    warnings
+    print
+    raise RuntimeError
+    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
 max-line-length = 120

setup.py

@@ -1,29 +1,57 @@
 #!/usr/bin/env python
 
+# Always prefer setuptools over distutils
 from setuptools import setup, find_packages
 
 # https://packaging.python.org/guides/single-sourcing-package-version/
 
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
+
+# 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 ...
+# the goal of the project is simplicity for researchers, don't want to add too much
+# engineer specific practices
 setup(
-    name="pytorch-lightning",
-    version='0.2.5.2',
-    description="The Keras for ML researchers using PyTorch",
-    author="William Falcon",
-    author_email="waf2107@columbia.edu",
-    url="https://github.com/williamFalcon/pytorch-lightning",
-    download_url="https://github.com/williamFalcon/pytorch-lightning",
-    license="MIT",
-    keywords=["deep learning", "pytorch", "AI"],
-    python_requires=">=3.5",
-    install_requires=[
-        "torch>=1.1.0",
-        "tqdm",
-        "test-tube>=0.6.7.1",
-    ],
+    name='pytorch-lightning',
+    version='0.5.1.2',
+    description='The Keras for ML researchers using PyTorch',
+    author='William Falcon',
+    author_email='waf2107@columbia.edu',
+    url='https://github.com/williamFalcon/pytorch-lightning',
+    download_url='https://github.com/williamFalcon/pytorch-lightning',
+    license='Apache-2',
     packages=find_packages(),
-    long_description=open("README.md", encoding="utf-8").read(),
+    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>=4.35.0',
+        'test-tube>=0.6.9',
+        'pandas>=0.20.3',
+    ],
+    classifiers=[
+        'Environment :: Console',
+        'Natural Language :: English',
+        # How mature is this project? Common values are
+        #   3 - Alpha, 4 - Beta, 5 - Production/Stable
+        'Development Status :: 4 - Beta',
+        # Indicate who your project is intended for
+        'Intended Audience :: Developers',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Scientific/Engineering :: Image Recognition',
+        'Topic :: Scientific/Engineering :: Information Analysis',
+        # Pick your license as you wish
+        'License :: OSI Approved :: BSD License',
+        'Operating System :: OS Independent',
+        # 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.6',
+        'Programming Language :: Python :: 3.7',
+    ],
 )

tests/README.md

@@ -0,0 +1,63 @@
+# PyTorch-Lightning Tests
+
+## Running tests
+The automatic travis tests ONLY run CPU-based tests. Although these cover most of the use cases,
+run on a 2-GPU machine to validate the full test-suite.
+
+
+To run all tests do the following:
+```bash
+git clone https://github.com/williamFalcon/pytorch-lightning
+cd pytorch-lightning
+
+# install module locally
+pip install -e .
+
+# install dev deps
+pip install -r requirements.txt
+
+# run tests
+py.test -v
+```
+
+To test models that require GPU make sure to run the above command on a GPU machine.
+The GPU machine must have:
+1. At least 2 GPUs.
+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. 
+
+```bash
+cd pytorch-lightning
+
+# generate coverage 
+pip install coverage
+coverage run --source pytorch_lightning -m py.test pytorch_lightning tests examples -v --doctest-modules
+
+# print coverage stats
+coverage report -m
+
+# exporting resulys
+coverage xml
+codecov -t 17327163-8cca-4a5d-86c8-ca5f2ef700bc  -v
+```
+
+

tests/debug.py

@@ -0,0 +1,114 @@
+from pytorch_lightning import Trainer
+from examples import LightningTemplateModel
+from pytorch_lightning.testing import LightningTestModel
+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 numpy as np
+import pdb
+from . import test_models
+
+
+class CoolModel(pl.LightningModule):
+
+    def __init(self):
+        super(CoolModel, self).__init__()
+        # not the best model...
+        self.l1 = torch.nn.Linear(28 * 28, 10)
+
+    def forward(self, x):
+        return torch.relu(self.l1(x))
+
+    def my_loss(self, y_hat, y):
+        return F.cross_entropy(y_hat, y)
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'training_loss': self.my_loss(y_hat, y)}
+
+    def validation_step(self, batch, batch_nb):
+        x, y = batch
+        y_hat = self.forward(x)
+        return {'val_loss': self.my_loss(y_hat, y)}
+
+    def validation_end(self, outputs):
+        avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
+        return avg_loss
+
+    def configure_optimizers(self):
+        return [torch.optim.Adam(self.parameters(), lr=0.02)]
+
+    @pl.data_loader
+    def train_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+
+    @pl.data_loader
+    def val_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+    @pl.data_loader
+    def test_dataloader(self):
+        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+
+def main():
+    """
+    Make sure DDP + AMP continue training correctly
+    :return:
+    """
+    """
+    Make sure DDP2 works
+    :return:
+    """
+    hparams = test_models.get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = test_models.init_save_dir()
+
+    # logger file to get meta
+    logger = test_models.get_test_tube_logger(False)
+    logger.log_hyperparams(hparams)
+    logger.save()
+
+    # logger file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    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 = test_models.load_model(logger.experiment, save_dir,
+                                              module_class=LightningTestModel)
+
+    new_trainer = Trainer(**trainer_options)
+    new_trainer.test(pretrained_model)
+
+    # test we have good test accuracy
+    test_models.assert_ok_test_acc(new_trainer)
+    test_models.clear_save_dir()
+
+
+if __name__ == '__main__':
+    main()

tests/requirements.txt

@@ -0,0 +1,9 @@
+tox
+coverage
+codecov
+pytest>=3.0.5
+pytest-cov
+flake8
+check-manifest
+test_tube
+mlflow