使用 pytorch-lightning

在用 PyTorch 训练不同模型时，需要编写很多重复的逻辑。一个工程的代码无可避免地变得越来越长，难以管理。PyTorch Lightning 作为一个对 PyTorch 二次封装的框架，能让训练逻辑的编写像堆积木一样秩序井然。

虽然叫做 lightning，这个库的学习成本并不低。好在一但熟悉，就能享受到很多便利。

安装

使用类似下面的指令安装 PyTorch Lightning。

python -m pip install lightning

完成后就可以在 Python 脚本里 import pytorch_lightning 了。

快速开始

PyTorch Lightning 的宗旨是：all in one。

不使用任何框架，写出来的训练代码会像是这样：

writer = SummaryWriter()
model = ResNet18().to(device)
optimizer = torch.optim.Adam(model.params, lr=1e-3)

for epoch in range(num_epochs):
    for i, batch in enumerate(train_data):
        x, y = batch
        x = x.to(device)
        output = model(x)
        loss = criterion(output, y)
        writer.add_scalar("train_loss", loss, num_epochs * len(train_data) + i)

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    with torch.no_grad():
        for batch in val_data:
            x, y = batch

                ···

若使用 PyTorch Lightning，可以把这一切封装到一个 class 内：

import pytorch_lightning as L

class LightningModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = ResNet18()
    
    def forward(self, x):
        return model(x)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters, lr=1e-3)
        return optimizer

    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        output = self.model(x)
        loss = criterion(output, y)
        self.log('train_loss', loss, on_step=True, on_epoch=False)
        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y = val_batch
        output = criterion(output, y)
        loss = criterion(output, y)
        self.log('val_loss', loss)

可见 lightning 的特点以及省下的工作：

• 将训练各步骤的代码分配到不同的函数，层次分明
• 不需要 .to(device)。lightning 会自动转换
• 不需要手动 loss.backward() optimizer.step() optimizer.zero_grad()
• validation 阶段不需要手动 torch.no_grad()
• 记录 log 时 lightning 会帮助填写当前 step 和 epoch
• ……

接下来训练的代码就可以非常简明：

model = LightningModel()
trainer = L.Trainer(accelerator='gpu', max_epochs=15)
trainer.fit(model, train_data, val_data) # 正式开始训练

使用混合精度

若要在 lightning 的训练启用混合精度加速，需要先执行：

torch.set_float32_matmul_precision('high')

然后在实例化 trainer 时，传入参数 precision='bf16-mixed' 即可。参数可选的值随硬件条件改变而变化，请查阅文档。

学习率调度器（scheduler）

在自定义 configure_optimizers() 时，除了需要指定 optimizer，还可以设定一个学习率调度器 scheduler。

class LightningModel(L.LightningModule):
    ···
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
        return [optimizer], [scheduler]

lightning 会自动执行 scheduler.step()。

trainer.init_module() 节省初始化资源

利用 trainer.init_module()，可以避免额外的精度转换，且能直接将模型加载到 GPU（而不是从 CPU 转移）。

trainer = Trainer(accelerator="cuda", precision="16-true")

with trainer.init_module():
    # models created here will be on GPU and in float16
    model = MyLightningModule()

若是从 checkpoint 初始化模型，可以向 trainer 传入参数 empty_init=True，这样在读取 checkpoint 之前模型的权重不会占用内存空间，且速度更快。

with trainer.init_module(empty_init=True):
    model = MyLightningModule.load_from_checkpoint("my/checkpoint/path.ckpt")

trainer.fit(model)

要注意，此时必须保证模型的每个权重都从 checkpoint 加载（或是手动加载），否则模型不完整。

针对使用 FSDP 或 DeepSpeed 训练的大参数模型，就不应使用 trainer.init_module() 了。对应的，为了加快大参数模型加载速度、减少内存消耗，在编写 LightningModel 时要把模型参数写到 def configure_model(self) 方法中。

设置所有随机数种子

lightning 提供了一种设置全局随机数种子的方法，能把 numpy、python 和 torch 的随机数种子都固定住，很适合调参。

若 dataloader 设置了多 worker，这个方法也能照顾到。保证每次重置获取的 batch 都一样。

import pytorch_lightning as pl
pl.seed_everything(42, workers=True)