Pytorch-lightning

Pytorch-lightning

简介

目前好像大多AI训练学习框架都使用的pytorch-lightning,因此今天也来了解一番，以后也要熟练使用，官方的定义为：构建和训练Pytorch 模型，并使用Lightning Apps模板将它们连接到ML 的生命周期，无需处理DIY基础设施，成本管理，扩展和其他令人头疼的问题。

• github地址：Lightning-AI/lightning

• 官方API 文档：Welcome to ⚡ PyTorch Lightning — PyTorch Lightning 1.8.0dev documentation )

How to Use

0. Install

pip install pytorch-lightning

1. Add the 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
   

2. Define a LightningModule (nn.Module)

   class LitAutoEncoder(pl.LightningModuel):
   	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):
   		embedding=self.encoder(x)
   		return embedding
   	
   	def training_step(self,batch,batch_idx):
   		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
   

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 feature

• 多GPU

  trainer=Trainer(max_epochs=1,accelerator='gpu',device=8)
  

• TPU

• 16 位精度

• 实验logging

• early_stopping

  es=EarlyStopping(monitor='val_loss')
  trainer=Trainer(callbacks=[checkpointing])
  

• model checkpoint

  checkpointing=ModelCheckpoint(monitor='val_loss')
  trainer=Trainer(callbacks=[checkpointing])
  

• torchscript

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

• ONNX

  # 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)
  

• training tricks
  40+的training trick供我们选择

Advantages

• 模型与硬件无关
• 代码简化
• 已于重构
• 犯更少的mistakes
• 保存了灵活性，但移除了大量样本
• 与流行的机器学习工具有集成
• 不同Python，Pytorch版本，操作系统，GPT进行支持
• 加快运行速度

手动控制训练过程

class LitAntoEncoder(pl.LightningModule):
	def __init__(self):
		super().__init__()
		self.automatic_optimization=False
		
	def training_step(self,batch,batch_idx):
		# access your optimizers with use_pl_optimizer=False. Default is True
        opt_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()

Example

Hello world

• MNIST

Contrastive Learning

• BYOL
• CPC v2
• Moco v2
• SIMCLR

NLP

• GPT-2
• BERT

Reinforcement Learning

• DQN
• Dueling-DQN
• Reinforce

Vision

• GAN

Classic ML

• Logistic Regression
• Linear Regression

官方API教程

Welcome to ⚡ PyTorch Lightning — PyTorch Lightning 1.8.0dev documentation (pytorch-lightning.readthedocs.io)

总结

Pytorch-lightning 作为2w star的github项目一定是很有用的，目前我仅仅尝试了一些example，需要完全掌握pytorch-ligthning中的简单语法，然后确实可以帮助我们减少重复AI代码的编写。