LysandreJik / pytorch-lightning

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

Home Page:

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool


PyTorch Lightning

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

MasterclassKey FeaturesHow To UseDocsExamplesCommunityLicence

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 3 steps guide

Trending contributors

Continuous Integration

System / PyTorch ver. 1.3 (min. req.)* 1.4 1.5 1.6 (latest)
Conda py3.7 [linux] PyTorch & Conda PyTorch & Conda PyTorch & Conda PyTorch & Conda
Linux py3.7 [GPUs**] - - - Build Status
Linux py3.7 [TPUs***] - - - TPU tests
Linux py3.6 / py3.7 / py3.8 CI complete testing - - CI complete testing
OSX py3.6 / py3.7 - CI complete testing - CI complete testing
Windows py3.6 / py3.7 / py3.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

How To Use

Step 0: Install

Simple installation from PyPI

pip install pytorch-lightning

From Conda

conda install pytorch-lightning -c conda-forge

Install bleeding-edge (no guarantees)

pip install git+ --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 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):
        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)
        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(), DataLoader(train), DataLoader(val))

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

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

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


# Distributes TPU core training
trainer = Trainer(tpu_cores=8)

# Single TPU core training
trainer = Trainer(tpu_cores=[1])

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,,, 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


Hello world

MNIST hello world

Contrastive Learning

CPC v2
Moco v2



Reinforcement Learning




Classic ML

Logistic Regression
Linear Regression


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.


Building open-source software with only a few part-time people is hard! We've secured funding to make sure we can hire a full-time staff, attend conferences, and move faster through implementing features you request.

Our goal is to build an incredible research platform and a big supportive community. Many open-source projects have gone on to fund operations through things like support and special help for big corporations!

If you are one of these corporations, please feel free to reach out to!


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


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

  title={PyTorch Lightning},
  author={Falcon, WA},
  journal={GitHub. Note:},
ezoic increase your site revenue


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

License:Apache License 2.0


Language:Python 94.9%Language:Jupyter Notebook 4.3%Language:Dockerfile 0.6%Language:Shell 0.2%Language:Jsonnet 0.0%