Build and train PyTorch models and connect them to the ML lifecycle using Lightning App templates, without handling DIY infrastructure, cost management, scaling, and other headaches.
Go to file
chaton ef8ef12fd0
[feat] pp 2/n (#5026)
* Added changes for RPC plugin

* Add missing kwargs

* Fix code format

* Loading refactors by introducing is_distributed var, fix optimizer step flow

* Add rpc guard

* Added docstrings and typing

* resolve comments

* Add additional rpc hook, refactor name of exit process hook for clarity

* remove annotation

* Modify behaviour to allow optional return, add test for rpc plugin

* resolve tests

* rename is_ddp_based

* update

* update for windows

* update

* resolve test

* code smell

* Added sequential plugin

* resolve bug

* update

* cleanup

* add Exception

* resolve docs

* Remove ddp support

* Revert distributed -> ddp

* Update pl_examples/basic_examples/conv_sequential_example.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

* Update pl_examples/basic_examples/conv_sequential_example.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

* Address code review points

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

* Add missing return

* Fix formatting, add datamodule args

* add small comment

* resolve comments

* resolve comments

* update source for fairscale

* update extras

* remove staticmethod

* resolve flake8

* Skip tests that are failing due to bug upstream with multiple optimizers and shard

* update

* update on comments

* clean test

* latest comments

* remove old comments

* add todo

* Update version

* update

* resolve bugs

* resolve bugs

* update test

* remove hanging test

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* resolve on comments

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* resolve on comments

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* Update pytorch_lightning/plugins/ddp_sequential_plugin.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

* remove ImportError

Co-authored-by: SeanNaren <sean@grid.ai>
Co-authored-by: Sean Naren <sean.narenthiran@gmail.com>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>
2020-12-09 12:56:51 +00:00
.circleci formatting (#4898) 2020-11-30 00:57:28 -05:00
.github fix ci: release (#5037) 2020-12-09 11:53:22 +01:00
benchmarks [feat] pp 2/n (#5026) 2020-12-09 12:56:51 +00:00
dockers Tpu save (#4309) 2020-12-02 13:05:11 +00:00
docs docs: minor spelling tweaks (#5022) 2020-12-08 16:27:43 -05:00
notebooks docs: default_root_path -> default_root_dir (#4942) 2020-12-02 19:17:34 -05:00
pl_examples [feat] pp 2/n (#5026) 2020-12-09 12:56:51 +00:00
pytorch_lightning [feat] pp 2/n (#5026) 2020-12-09 12:56:51 +00:00
requirements [feat] pp 2/n (#5026) 2020-12-09 12:56:51 +00:00
tests [feat] pp 2/n (#5026) 2020-12-09 12:56:51 +00:00
.codecov.yml skip files in coverage (#3944) 2020-10-07 12:37:01 -04:00
.drone.jsonnet Create .drone.jsonnet (#4968) 2020-12-06 00:17:18 +00:00
.drone.yml Simplify optimization Logic (#4984) 2020-12-07 12:55:49 +00:00
.gitignore [FEAT] Add lambda closure to manual_optimizer_step (#4618) 2020-11-12 19:22:06 +00:00
.mergify.yml run TPU tests with multiple versions (#3024) 2020-09-30 08:36:02 -04:00
.pep8speaks.yml Set pep8speaks' max-line-length to 120 (same as black) (#3173) 2020-08-25 21:21:02 -04:00
.pre-commit-config.yaml [feat] pp 2/n (#5026) 2020-12-09 12:56:51 +00:00
.readthedocs.yml move base req. to root (#4219) 2020-10-18 20:40:18 +02:00
.update.sh default test logger (#1478) 2020-04-21 20:33:10 -04:00
CHANGELOG.md adding missing changelogs (#5019) 2020-12-09 11:53:37 +00:00
LICENSE update license (#809) 2020-02-09 14:18:50 -05:00
MANIFEST.in CI: update badges for release (#5002) 2020-12-09 10:59:44 +01:00
Makefile [make] Create Makefile (#4620) 2020-11-12 09:25:31 +00:00
README.md CI: update badges for release (#5002) 2020-12-09 10:59:44 +01:00
environment.yml lock pytorch nightly version (#4469) 2020-11-01 23:47:51 +05:30
pyproject.toml Apply import formatting to files in the 2nd top level (#4717) 2020-11-18 00:29:09 +01:00
requirements.txt upgrade min deps (#4934) 2020-12-01 17:19:44 +00:00
setup.cfg replace pyright by mypy (#5021) 2020-12-09 10:57:11 +08:00
setup.py Added the function for downloading the badges locally and replace the url with downlaod path (#4250) 2020-12-01 22:58:32 +01:00

README.md

The lightweight PyTorch wrapper for high-performance AI research.
Scale your models, not the boilerplate.


WebsiteKey FeaturesHow To UseDocsExamplesCommunityGrid AILicence

PyPI - Python Version PyPI Status PyPI Status Conda DockerHub codecov

ReadTheDocs Slack Discourse status license Next Release

*Codecov is > 90%+ but build delays may show less

PyTorch Lightning is just organized PyTorch

Lightning disentangles PyTorch code to decouple the science from the engineering. PT to PL


Lightning Philosophy

Lightning is designed with these principles in mind:

Principle 1: Enable maximal flexibility.
Principle 2: Abstract away unecessary boilerplate, but make it accessible when needed.
Principle 3: Systems should be self-contained (ie: optimizers, computation code, etc).
Principle 4: Deep learning code should be organized into 4 distinct categories.

  • Research code (the LightningModule).
  • Engineering code (you delete, and is handled by the Trainer).
  • Non-essential research code (logging, etc... this goes in Callbacks).
  • Data (use PyTorch Dataloaders or organize them into a LightningDataModule).

Once you do this, you can train on multiple-GPUs, TPUs, CPUs and even in 16-bit precision without changing your code!

Get started with our 2 step guide


Inference

Lightning is also designed for the fast inference AI researchers and production teams need to scale up things like BERT and self-supervised learning. Lightning can automatically export to ONNX or TorchScript for those cases.



Continuous Integration

System / PyTorch ver. 1.3 (min. req.)* 1.4 1.5 1.6 1.7 (latest) 1.8 (nightly)
Conda py3.7 [linux] PyTorch & Conda PyTorch & Conda PyTorch & Conda PyTorch & Conda PyTorch & Conda PyTorch & Conda
Linux py3.7 [GPUs**] - - - GPUs Status - -
Linux py3.{6,7} [TPUs***] - - - TPU tests TPU tests -
Linux py3.{6,7} CI complete testing - - - CI complete testing -
OSX py3.{6,7,8} - CI complete testing - - CI complete testing -
Windows py3.{6,7,8} CI complete testing - - - CI complete testing -
  • * torch>=1.4 is the minimal pytorch version for Python 3.8
  • ** tests run on two NVIDIA K80
  • *** tests run on Google GKE TPUv2/3
  • TPU w/ py3.6/py3.7 means we support Colab and Kaggle env.

How To Use

Step 0: Install

Simple installation from PyPI

pip install pytorch-lightning

To get full package experience you can install also all optional dependencies with pytorch-lightning['extra'] or for CPU users with pytorch-lightning['cpu-extra'].

From Conda

conda install pytorch-lightning -c conda-forge

Install bleeding-edge (no guarantees)

pip install git+https://github.com/PytorchLightning/pytorch-lightning.git@master --upgrade

Step 0: Add these imports

import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
import pytorch_lightning as pl

Step 1: Define a LightningModule (nn.Module subclass)

A LightningModule defines a full system (ie: a GAN, autoencoder, BERT or a simple Image Classifier).

class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))
    
    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step defined the train loop. It is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer
Note: Training_step defines the training loop. Forward defines how the LightningModule behaves during inference/prediction.

Step 2: Train!

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train, val = random_split(dataset, [55000, 5000])

autoencoder = LitAutoEncoder()
trainer = pl.Trainer()
trainer.fit(autoencoder, DataLoader(train), DataLoader(val))

And without changing a single line of code, you could run on GPUs/TPUs

# 8 GPUs
trainer = Trainer(max_epochs=1, gpus=8)

# 256 GPUs
trainer = Trainer(max_epochs=1, gpus=8, num_nodes=32)

# TPUs
trainer = Trainer(tpu_cores=8)

And even export for production via onnx or torchscript

# torchscript
autoencoder = LitAutoEncoder()
torch.jit.save(autoencoder.to_torchscript(), "model.pt") 

# onnx
with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
    autoencoder = LitAutoEncoder()
    input_sample = torch.randn((1, 64))
    autoencoder.to_onnx(tmpfile.name, input_sample, export_params=True)
    os.path.isfile(tmpfile.name)

For advanced users, you can still own complex training loops

class LitAutoEncoder(pl.LightningModule):
    def training_step(self, batch, batch_idx, opt_idx):
        (opt_a, opt_b) = self.optimizers()
        
        loss_a = ...
        self.manual_backward(loss_a, opt_a)
        opt_a.step()
        opt_a.zero_grad()
        
        loss_b = ...
        self.manual_backward(loss_b, opt_b, retain_graph=True)
        self.manual_backward(loss_b, opt_b)
        opt_b.step()
        opt_b.zero_grad()

Key Features

  • Scale your models to run on any hardware (CPU, GPUs, TPUs) without changing your model
  • Making code more readable by decoupling the research code from the engineering
  • Easier to reproduce
  • Less error prone by automating most of the training loop and tricky engineering
  • Keeps all the flexibility (LightningModules are still PyTorch modules), but removes a ton of boilerplate
  • Lightning has out-of-the-box integration with the popular logging/visualizing frameworks (Tensorboard, MLFlow, Neptune.ai, Comet.ml, Wandb).
  • Tested rigorously with every new PR. We test every combination of PyTorch and Python supported versions, every OS, multi GPUs and even TPUs.
  • Minimal running speed overhead (about 300 ms per epoch compared with pure PyTorch).

Lightning automates 40+ parts of DL/ML research

  • GPU training
  • Distributed GPU (cluster) training
  • TPU training
  • EarlyStopping
  • Logging/Visualizing
  • Checkpointing
  • Experiment management
  • Full list here

Examples

Hello world

MNIST hello world
MNIST on TPUs

Contrastive Learning

BYOL
CPC v2
Moco v2
SIMCLR

NLP

BERT
GPT-2

Reinforcement Learning

DQN
Dueling-DQN
Reinforce

Vision

GAN

Classic ML

Logistic Regression
Linear Regression


Community

The lightning community is maintained by

  • 16 core contributors who are all a mix of professional engineers, Research Scientists, Ph.D. students from top AI labs.
  • 280+ community contributors.

Lightning is also part of the PyTorch ecosystem which requires projects to have solid testing, documentation and support.

Asking for help

If you have any questions please:

  1. Read the docs.
  2. Look it up in our forum (or add a new question)
  3. Search through the issues.
  4. Join our slack.
  5. Ask on stackoverflow with the tag pytorch-lightning.

Funding

Building open-source software with only a few part-time people is hard!

We're venture funded and backed by some of the top VC funds in the world, Index Ventures, Bain Capital Ventures, First Minute Capital.

Their funding ensures we can continue to build awesome tooling like Grid, give you around the clock support, hire a full-time staff, attend conferences, and move faster through implementing features you request.

To supercharge your research and production work, visit our Grid.ai platform


Grid AI

Grid AI is our native platform for training models at scale on the cloud!

Sign up for early access here

To use grid, take your regular command:

    python my_model.py --learning_rate 1e-6 --layers 2 --gpus 4

And change it to use the grid train command:

    grid train --grid_gpus 4 my_model.py --learning_rate 'uniform(1e-6, 1e-1, 20)' --layers '[2, 4, 8, 16]'

The above command will launch (20 * 4) experiments each running on 4 GPUs (320 GPUs!) - by making ZERO changes to your code.


Licence

Please observe the Apache 2.0 license that is listed in this repository. In addition the Lightning framework is Patent Pending.

BibTeX

If you want to cite the framework feel free to use this (but only if you loved it 😊):

@article{falcon2019pytorch,
  title={PyTorch Lightning},
  author={Falcon, WA},
  journal={GitHub. Note: https://github.com/PyTorchLightning/pytorch-lightning},
  volume={3},
  year={2019}
}