代码拉取完成，页面将自动刷新
The lightweight PyTorch wrapper for high-performance AI research. Scale your models, not the boilerplate.
Website • Key Features • How To Use • Docs • Examples • Community • Grid AI • Licence
Lightning disentangles PyTorch code to decouple the science from the engineering.
Lightning is designed with these principles in mind:
Principle 1: Enable maximal flexibility. Principle 2: Abstract away unnecessary 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.
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
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.
| System / PyTorch ver. | 1.3 (min. req.)* | 1.4 | 1.5 | 1.6 | 1.7 (latest) | 1.8 (nightly) |
|---|---|---|---|---|---|---|
| Conda py3.7 [linux] | |
|
|
|
|
|
| Linux py3.7 [GPUs**] | - | - | - | |
- | - |
| Linux py3.{6,7} [TPUs***] | - | - | - | |
|
- |
| Linux py3.{6,7} | |
- | - | - | |
- |
| OSX py3.{6,7,8} | - | |
- | - | |
- |
| Windows py3.{6,7,8} | |
- | - | - | |
- |
torch>=1.4 is the minimal pytorch version for Python 3.8Simple 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
the actual status of 1.2 [nightly] is following:
Install future release from the source (no guarantees)
pip install git+https://github.com/PytorchLightning/pytorch-lightning.git@release/1.2-dev --upgrade
or nightly from testing PyPI
pip install -iU https://test.pypi.org/simple/ pytorch-lightning
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
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.
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))
# 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)
# 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)
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()
The lightning community is maintained by
Lightning is also part of the PyTorch ecosystem which requires projects to have solid testing, documentation and support.
If you have any questions please:
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 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.
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 😊):
@article{falcon2019pytorch,
title={PyTorch Lightning},
author={Falcon, WA},
journal={GitHub. Note: https://github.com/PyTorchLightning/pytorch-lightning},
volume={3},
year={2019}
}
此处可能存在不合适展示的内容，页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容，可点击提交进行申诉，我们将尽快为您处理。
马建仓 AI 助手