Pytorch-lightning
简介
目前好像大多AI训练学习框架都使用的pytorch-lightning,因此今天也来了解一番,以后也要熟练使用,官方的定义为:构建和训练Pytorch 模型,并使用Lightning Apps模板将它们连接到ML 的生命周期,无需处理DIY基础设施,成本管理,扩展和其他令人头疼的问题。
- github地址:Lightning-AI/lightning
How to Use
- Install
pip install pytorch-lightning
-
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
-
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
-
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教程
总结
Pytorch-lightning 作为2w star的github项目一定是很有用的,目前我仅仅尝试了一些example,需要完全掌握pytorch-ligthning中的简单语法,然后确实可以帮助我们减少重复AI代码的编写。