Lightning in 15 minutes#

Lightning in 15 minutes — PyTorch Lightning 2.0.4 documentation

安装 PyTorch Lightning#

pip install lightning

conda install lightning -c conda-forge

定义一个LightningModule#

LightningModule可以让pytorch的nn.Module可以整合一些训练过程（也可以有验证和测试）。

如下是一个手写数字识别自动编码器（autoencoder）的样例：

1
import os
2
import torch
3
from torch import optim, nn, utils
4
from torchvision.datasets import MNIST
5
from torchvision.transforms import ToTensor
6
import lightning.pytorch as pl
7

8
'''
9
定义两个模型，编码器和解码器，这个是pytorch的模型对象
10
'''
11
encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
12
decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))
13

14
# 定义LightningModule
15
class LitAutoEncoder(pl.LightningModule):
16
    def __init__(self, encoder, decoder):
17
        super().__init__()
18
        self.encoder = encoder
19
        self.decoder = decoder
20

21
    def training_step(self, batch, batch_idx):
22
        # 训练步骤
23
        # 这个跟 forward 不相关
24
        x, y = batch
25
        x = x.view(x.size(0), -1)
26
        z = self.encoder(x)
27
        x_hat = self.decoder(z)
28
        loss = nn.functional.mse_loss(x_hat, x)
29
        # 存储日志（需要安装Tensorboard）
30
        self.log("train_loss", loss)
31
        return loss
32

33
    def configure_optimizers(self):
34
        # 优化器
35
        optimizer = optim.Adam(self.parameters(), lr=1e-3)
36
        return optimizer
37

38
# 初始化自动编码器
39
autoencoder = LitAutoEncoder(encoder, decoder)

定义数据集#

Lightning支持所有可迭代的数据集形式（DataLoader，numpy，以及其他）。

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

训练模型#

Lightning的Trainer对象可以整合LightningModule与不同数据集，并扩展了一些工程所需方法。

1
# 训练模型
2
trainer = pl.Trainer(limit_train_batches=100, max_epochs=10)
3
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

Trainer对象也实现了很多常用的过程：

Epoch和batch迭代。
optimizer.step(), loss.backward(), optimizer.zero_grad()
验证过程中的**model.eval()。**
模型存储和载入
Tensorboard
多GPU
TPU
半精度混合

【注意】：在jupyter下，多卡训练可能会报错，可以试试直接用python代码。

使用模型#

训练完模型后，可以导出到 onnx、torchscript 并将其投入生产，或者只是加载权重并运行预测。

1
# 载入模型
2
checkpoint = "./lightning_logs/version_0/checkpoints/epoch=0-step=100.ckpt"
3
autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint, encoder=encoder, decoder=decoder)
4

5
# 选择训练好的编码器
6
encoder = autoencoder.encoder
7
encoder.eval()
8

9
# 编码图片
10
fake_image_batch = torch.randn(8, 28 * 28).to(next(encoder.parameters()).device)
11
embeddings = encoder(fake_image_batch)
12
print("⚡" * 20, "\nPredictions (4 image embeddings):\n", embeddings, "\n", "⚡" * 20)

训练可视化#

如果安装了Tensorboard，可以用它来观察实验过程。

1
tensorboard --logdir .

额外训练设置#

1
# 4gpu训练
2
trainer = Trainer(
3
    devices=4,
4
    accelerator="gpu",
5
 )
6

7
# train 1TB+ parameter models with Deepspeed/fsdp
8
# 使用 Deepspeed 训练大模型
9
trainer = Trainer(
10
    devices=4,
11
    accelerator="gpu",
12
    strategy="deepspeed_stage_2",
13
    precision=16
14
 )
15

16
# 20+ helpful flags for rapid idea iteration
17
# 有助于快速迭代的一些设置
18
trainer = Trainer(
19
    max_epochs=10,
20
    min_epochs=5,
21
    overfit_batches=1
22
 )
23

24
# access the latest state of the art techniques
25
# 获取最新的技术
26
trainer = Trainer(callbacks=[StochasticWeightAveraging(...)])

一些灵活设置#

定制训练循环#

LightningModule中设置了20多种断点（HOOK），可以用来定制训练过程：

1
class LitAutoEncoder(pl.LightningModule):
2
    def backward(self, loss):
3
        loss.backward()

扩展Trainer#

在上面这个代码种，对模型的存储进行了一些设置。这些设置可以在pl.Callback对象中实现，并导入Trainer对象。

用如下方式可以导入Trainer:

1
trainer = Trainer(callbacks=[AWSCheckpoints()])

原文链接：CSDN