Package: pytorch-lightning
The lightweight PyTorch wrapper for high-performance AI research.
Scale your models, not the boilerplate.Website•Key Features•How To Use•Docs•Examples•Community•Lightning AI•License*Codecov is > 90%+ but build delays may show lessPyTorch Lightning is just organized PyTorchLightning disentangles PyTorch code to decouple the science from the engineering.Lightning Design PhilosophyLightning structures PyTorch code with these principles:Lightning forces the following structure to your code which makes it reusable and shareable: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, IPUs, HPUs and even in 16-bit precision without changing your code!Get started in just 15 minutesContinuous IntegrationLightning is rigorously tested across multiple CPUs, GPUs and TPUs and against major Python and PyTorch versions.Current build statusesSystem / PyTorch ver.1.121.132.02.1Linux py3.9 [GPUs]Linux py3.9 [TPUs]Linux (multiple Python versions)OSX (multiple Python versions)Windows (multiple Python versions)How To UseStep 0: InstallSimple installation from PyPIpipinstallpytorch-lightningStep 1: Add these importsimportosimporttorchfromtorchimportnnimporttorch.nn.functionalasFfromtorchvision.datasetsimportMNISTfromtorch.utils.dataimportDataLoader,random_splitfromtorchvisionimporttransformsimportpytorch_lightningasplStep 2: Define a LightningModule (nn.Module subclass)A LightningModule defines a fullsystem(ie: a GAN, autoencoder, BERT or a simple Image Classifier).classLitAutoEncoder(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))defforward(self,x):# in lightning, forward defines the prediction/inference actionsembedding=self.encoder(x)returnembeddingdeftraining_step(self,batch,batch_idx):# training_step defines the train loop. It is independent of forwardx,y=batchx=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)returnlossdefconfigure_optimizers(self):optimizer=torch.optim.Adam(self.parameters(),lr=1e-3)returnoptimizerNote: Training_step defines the training loop. Forward defines how the LightningModule behaves during inference/prediction.Step 3: 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))Advanced featuresLightning has over40+ advanced featuresdesigned for professional AI research at scale.Here are some examples:Highlighted feature code snippets# 8 GPUs# no code changes neededtrainer=Trainer(max_epochs=1,accelerator="gpu",devices=8)# 256 GPUstrainer=Trainer(max_epochs=1,accelerator="gpu",devices=8,num_nodes=32)Train on TPUs without code changes# no code changes neededtrainer=Trainer(accelerator="tpu",devices=8)16-bit precision# no code changes neededtrainer=Trainer(precision=16)Experiment managersfrompytorch_lightningimportloggers# tensorboardtrainer=Trainer(logger=TensorBoardLogger("logs/"))# weights and biasestrainer=Trainer(logger=loggers.WandbLogger())# comettrainer=Trainer(logger=loggers.CometLogger())# mlflowtrainer=Trainer(logger=loggers.MLFlowLogger())# neptunetrainer=Trainer(logger=loggers.NeptuneLogger())# ... and dozens moreEarlyStoppinges=EarlyStopping(monitor="val_loss")trainer=Trainer(callbacks=[es])Checkpointingcheckpointing=ModelCheckpoint(monitor="val_loss")trainer=Trainer(callbacks=[checkpointing])Export to torchscript (JIT) (production use)# torchscriptautoencoder=LitAutoEncoder()torch.jit.save(autoencoder.to_torchscript(),"model.pt")Export to ONNX (production use)autoencoder=LitAutoEncoder()input_sample=torch.randn((1,64))withtempfile.NamedTemporaryFile(suffix=".onnx",delete=False)astmpfile:autoencoder.to_onnx(tmpfile.name,input_sample,export_params=True)Pro-level control of optimization (advanced users)For complex/professional level work, you have optional full control of the optimizers.classLitAutoEncoder(pl.LightningModule):def__init__(self):super().__init__()self.automatic_optimization=Falsedeftraining_step(self,batch,batch_idx):# access your optimizers with use_pl_optimizer=False. Default is Trueopt_a,opt_b=self.optimizers(use_pl_optimizer=True)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()Advantages over unstructured PyTorchModels become hardware agnosticCode is clear to read because engineering code is abstracted awayEasier to reproduceMake fewer mistakes because lightning handles the tricky engineeringKeeps all the flexibility (LightningModules are still PyTorch modules), but removes a ton of boilerplateLightning has dozens of integrations with popular machine learning tools.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).ExamplesSelf-supervised LearningCPC transformsMoco v2 transformsSimCLR transformsConvolutional ArchitecturesGPT-2UNetReinforcement LearningDQN LossDouble DQN LossPer DQN LossGANsBasic GANDCGANClassic MLLogistic RegressionLinear RegressionCommunityThe PyTorch Lightning community is maintained by10+ core contributorswho are all a mix of professional engineers, Research Scientists, and Ph.D. students from top AI labs.680+ active community contributors.Want to help us build Lightning and reduce boilerplate for thousands of researchers?Learn how to make your first contribution herePyTorch Lightning is also part of thePyTorch ecosystemwhich requires projects to have solid testing, documentation and support.Asking for helpIf you have any questions please:Read the docs.Search through existing Discussions, oradd a new questionJoin our Discord community.