【PyTorch Lightning】简介

目录

一、Lightning 简约哲学

1.1 研究代码 (Research code)

1.2 工程代码 (Engineering code)

1.3 非必要代码 (Non-essential code)

二、典型的 AI 研究项目

三、生命周期

四、使用 Lightning 的好处

---

相关文章

【PyTorch Ligntning】快速上手简明指南 

【PyTorch Ligntning】如何将 PyTorch 组织为 Lightning

【PyTorch Lightning】1.0 正式发布：从 0 到 1 

---

PyTorch 已足够简单易用，但简单易用不等于方便快捷。特别是做大量实验时，很多东西都会变得复杂，代码也会变得庞大，此时就容易出错。

针对该问题，就有了 PyTorch Lightning。它可以重构你的 PyTorch 代码，抽出复杂重复部分，让你专注于核心的构建，让你的实验更快速更便捷地开展迭代。

---

一、Lightning 简约哲学

大部分的 DL/ML 代码都可以分为以下这三部分：

• 研究代码 Research code
• 工程代码 Engineering code
• 非必要代码 Non-essential code

1.1 研究代码 (Research code)

这部分属于模型（神经网络）部分，一般处理模型的结构、训练等定制化部分。

在 Linghtning 中，这部分代码抽象为 LightningModule 类。

1.2 工程代码 (Engineering code)

这部分代码很重要的特点是：重复性强，如设置 early stopping、16位精度、GPUs 分布训练等。

在 Linghtning 中，这部分抽象为 Trainer 类。

1.3 非必要代码 (Non-essential code)

这部分代码有利于实验的进行，但和实验没有直接关系，甚至可以不用。如检查梯度、向 tensorboard 输出 log。

在 Linghtning 中，这部分抽象为 Callbacks 类。

---

二、典型的 AI 研究项目

在大多数研究项目中，研究代码 通常可以归纳到以下关键部分：

• 模型
• 训练/验证/测试 数据
• 优化器
• 训练/验证/测试 计算

上面已经提到，研究代码 在 Lightning 中，是抽象为 LightningModule 类；而该类与我们平时在 PyTorch 中使用的 torch.nn.Module 是一样的 (在原有代码中直接替换 Module 而不改其他代码也可以)。但不同的是，Lightning 围绕 torch.nn.Module 做了很多功能性的补充，把上面 4 个关键部分都囊括了进来。

如此设定的意义在于：我们的 研究代码 都是围绕 神经网络模型 来运行的，所以 Lightning 把这部分代码都集合在一个类里。所以接下来的介绍，都围绕 LightningModule 类来展开。

---

三、生命周期

为先呈现一个总体的概念，此处先介绍 LightningModule 中运行的生命流程。

以下所有函数都在 LightningModule 类中。

这部分是训练开始之后的执行 “一般（默认）顺序”。

• 首先是准备工作，包括初始化 LightningModule，准备数据 和 配置优化器。

这部分代码 只执行一次。

1. `__init__()`（初始化 LightningModule ）
2. `prepare_data()` （准备数据，包括下载数据、预处理等等）
3. `configure_optimizers()` （配置优化器）

• 测试 “验证代码”。

提前来做的意义在于：不需要等待漫长的训练过程才发现验证代码有错。

这部分就是提前执行 “验证代码”，所以和下面的验证部分是一样的。

1. `val_dataloader()`
2. `validation_step()`
3. `validation_epoch_end()`

• 开始加载dataloader，用来给训练加载数据

1. `train_dataloader()` 
2. `val_dataloader()` （如果你定义了）

• 下面部分就是循环训练了，_step() 指按 batch 进行的部分；_epoch_step() 指所有 batch 执行完后 (一个 epoch) 要进行的部分。

# 循环训练与验证
1. `training_step()`
2. `validation_step()`
3. `validation_epoch_end()`

• 最后训练完了，就要进行测试，但测试部分需手动调用 .test()，以避免误操作。

# 测试（需要手动调用）
1. `test_dataloader()` 
2. `test_step()`
3. `test_epoch_end()`

不难总结，在训练部分，主要包含三部分：_dataloader / _step / _epoch_end。Lightning 把训练的三部分抽象成三个函数，而用户只需要“填鸭式”地补充这三部分，就可以完成模型训练部分代码的编写。

为更清晰地展现这三部分的具体位置，以下用 PyTorch 实现方式 来展现其位置。

for epoch in epochs:
    for batch in train_dataloader:
        # train_step
        # ....
        # train_step
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    for batch in val_dataloader:
        # validation_step
        # ....
        # validation_step
    
    # *_step_end
    # ....
    # *_step_end

---

四、使用 Lightning 的好处

• 只需专注于 研究代码

不需要写一大堆的 .cuda() 和 .to(device)，Lightning 会自动处理。如果要新建一个 tensor，可以使用 type_as 来使得新tensor 处于相同的处理器上。

def training_step(self, batch, batch_idx):
    x, y = batch

    # 把z放在和x一样的处理器上
    z = sample_noise()
    z = z.type_as(x)

此处需要注意的是，不是所有的在 LightningModule 的 tensor 都会被自动处理，而是只有从 Dataloader 里获取的 tensor 才会被自动处理，所以对于 transductive learning 的训练，最好自己写 Dataloader 的处理函数。 

• 工程代码参数化

平时写模型训练时，这部分代码会不断重复，但又不得不做，比如 early stopping、精度调整、显存内存间数据转移等。这部分代码虽然不难，但减少这部分代码会使得 研究代码 更加清晰，整体也更加简洁。

下面是简单的展示，表示使用 LightningModule 建立好模型后，如何进行训练。

model = LightningModuleClass()
trainer = pl.Trainer(gpus="0",  # 用来配置使用什么GPU
                     precision=32, # 用来配置使用什么精度，默认是32
                     max_epochs=200 # 迭代次数
                     )

trainer.fit(model)  # 开始训练
trainer.test()  # 训练完之后测试

---

参考文献

https://zhuanlan.zhihu.com/p/120331610

https://pytorch-lightning.readthedocs.io/en/latest/