PyTorch Lightning(一)、LIGHTNING IN 15 MINUTES

LIGHTNING IN 15 MINUTES

Lightning in 15 minutes — PyTorch Lightning 2.2.5 documentation

背景知识：不需要

目标：了解lightning工作流的7个关键步骤

将pytorch转lightning主要了的部分

1. 模型结构放入__init__

   

2. 将编码器放入forward hook

   

3. 配置优化器，参数本省就包含来自编码器和解码器。

   

4. 创建训练

   

5. 设置验证环节

   

6. 把有关于cuda的部分移除掉，后面会统一追加

   

7. 设置反向传播和优化器部分，也可以不显式的写明

   

8. 模型初始化

   

9. 初始化训练器

   

   优化器会自动的组合这些步骤

   

10. 设置数据载入

    

11. 设置回调函数

    

12. 快速的设置GPU或者CPU

创建lightning的7个主要步骤

1. 安装

   pip install lightning
   

2. 定义LightningModel，当然validation_step和test_step也是一样的

   import os
   from torch import optim, nn, utils, Tensor
   from torchvision.datasets import MNIST
   from torchvision.transforms import ToTensor
   import lightning as L
   
   # define any number of nn.Modules (or use your current ones)
   encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
   decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))
   
   # define the LightningModule
   class LitAutoEncoder(L.LightningModule):
       def __init__(self, encoder, decoder):
           super().__init__()
           self.encoder = encoder
           self.decoder = decoder
   
       def training_step(self, batch, batch_idx):
           # training_step defines 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 = nn.functional.mse_loss(x_hat, x)
           # Logging to TensorBoard (if installed) by default
           self.log("train_loss", loss)
           return loss
   
       def configure_optimizers(self):
           optimizer = optim.Adam(self.parameters(), lr=1e-3)
           return optimizer
   
   # init the autoencoder
   autoencoder = LitAutoEncoder(encoder, decoder)
   

3. 定义一个dataload

   # setup data
   dataset = MNIST(os.getcwd(), download=True, transform=ToTensor())
   train_loader = utils.data.DataLoader(dataset)
   

4. 开启train

   # train the model (hint: here are some helpful Trainer arguments for rapid idea iteration)
   trainer = L.Trainer(limit_train_batches=100, max_epochs=1)
   trainer.fit(model=autoencoder, train_dataloaders=train_loader)
   

   trainer里面有很多的小技巧

   • Epoch and batch iteration
   • optimizer.step(), loss.backward(), optimizer.zero_grad() calls
   • Calling of model.eval(), enabling/disabling grads during evaluation
   • Checkpoint Saving and Loading
   • Tensorboard (see loggers options)
   • Multi-GPU support
   • TPU
   • 16-bit precision AMP support

5. 模型预测

   # load checkpoint
   checkpoint = "./lightning_logs/version_0/checkpoints/epoch=0-step=100.ckpt"
   autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint, encoder=encoder, decoder=decoder)
   
   # choose your trained nn.Module
   encoder = autoencoder.encoder
   encoder.eval()
   
   # embed 4 fake images!
   fake_image_batch = torch.rand(4, 28 * 28, device=autoencoder.device)
   embeddings = encoder(fake_image_batch)
   print("⚡" * 20, "\nPredictions (4 image embeddings):\n", embeddings, "\n", "⚡" * 20)
   

6. 训练可视化，如果有tensorboard的话，直接打开网页http://localhost:6006/

   tensorboard --logdir .
   

7. 改变训练参数

   # train on 4 GPUs
   trainer = L.Trainer(
       devices=4,
       accelerator="gpu",
    )
   
   # train 1TB+ parameter models with Deepspeed/fsdp
   trainer = L.Trainer(
       devices=4,
       accelerator="gpu",
       strategy="deepspeed_stage_2",
       precision=16
    )
   
   # 20+ helpful flags for rapid idea iteration
   trainer = L.Trainer(
       max_epochs=10,
       min_epochs=5,
       overfit_batches=1
    )
   
   # access the latest state of the art techniques
   trainer = L.Trainer(callbacks=[StochasticWeightAveraging(...)])
   

其他使用技巧

1. 自定义的train loop

   

2. Trainer的扩展

   

   如果您有多行具有相似功能的代码，您可以使用回调将它们轻松地分组在一起，并同时打开或关闭所有这些行。

3. 使用pytorch的原始循环

4. 手动构建优化器