pytorch lightning最简上手

pytorch lightning最简上手

pytorch lightning 是对原生 pytorch 的通用模型开发过程进行封装的一个工具库。本文不会介绍它的高级功能，而是通过几个最简单的例子来帮助读者快速理解、上手基本的使用方式。在掌握基础 API 和使用方式之后，读者可自行到 pytorch lightning 的官方文档，了解进阶 API。本文假设读者对原生 pytorch 训练脚本的搭建方法已经比较熟悉。

安装

pytorch lighning 的安装非常简单，直接使用 pip 安装即可：

pip install pytorch-lightning
1

最简例子

pytorch lightning 有两个最核心的 API：LigtningModule 和 Trainer 。

其中 LightningModule 是我们熟悉的 torch.nn.Module 的子类，可以通过

print(isinstance(pl.LightningModule(), torch.nn.Module))
1

来验证。这意味着该类同样需要实现 forward 方法，并可直接通过实例调用。

Trainer 则是开始执行模型训练、测试过程的类，传入一个 LightningModule 和对应控制参数来实例化即可开始训练。

我们从一个最简单的例子——MNIST 手写数字识别开始：

1 导入必要的库

导入 pytorch_lightning 和 pytorch 常用的库。

import os

import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
import pytorch_lightning as pl
1
2
3
4
5
6
7
8

2 实现最简LigntningModule

我们先实现一个最简的 LightningModule。

• __init__

  构造函数中，像常见的 torch.nn.Module 一样，我们定义好模型的层。由于是最简实例，这里只有一层线性层，将手写数字图像映射为输出 logits。

• forward

  由于是继承自 torch.nn.Module，因此实现 forward 方法是必须的。forward 方法要完成模型的前向过程，这里直接调用 __init__ 中定义好的线性层，完成模型前向过程。

• train_dataloader

  train_dataloader 方法也是最简实现中必须的，它的功能是获取训练集的 DataLoader。这里我们返回 MNIST 数据集的 DataLoader。dataloader 的获取也可以不在类内实现，而是在 fit 时传入，后面会介绍。

• training_step

  training_step 是是 LigtningModule 的核心方法，它定义了一个训练步中需要做的事情。在深度学习的训练步中，最核心的事情就是模型前向，得到结果，计算损失，反向传播，更新参数，这几步在 pytorch 中都有对应的方法供调用。但是在 pytorch lightning 中，我们只需要进行模型前向，并返回必要的信息即可。在最简实现中，我们只需返回损失。

• configure_optimizer

  在 training_step 中，我们只需返回损失，这意味着模型的反向传播和参数更新过程由 pytorch lightning 帮我们完成了。虽然这个过程可以有框架自己完成，但是我们还是要指定参数更新所用的优化器，在很多模型中，优化器、学习率等超参数设置对结果影响很大。在最简实现中，我们设置好学习率，并返回一个 Adam 优化器。

class MNISTModel(pl.LightningModule):

    def __init__(self):
        super(MNISTModel, self).__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))
      
    def train_dataloader(self):
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
      
    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

以上我们实现 training_step，train_dataloader, configure_optimizer，已经是最简单的 LightningModule 的实现了。如果连这三个方法都没有实现的话，将会报错：

 No `xxx` method defined. Lightning `Trainer` expects as minimum a `training_step()`, `train_dataloader()` and `configure_optimizers()` to be defined
1

3 开始训练

在实现好 LightningModule 之后，就可以开始训练了。

启动训练的最简实现非常简单，只需三行：实例化模型、实例化训练器、开始训练！

model = MNISTModel()
trainer = pl.Trainer(gpus=1, max_epochs=2)
trainer.fit(model)
1
2
3

开始训练后，pytorch lightning 会打印出可用设备、模型参数等丰富的信息。

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2]

  | Name | Type   | Params
--------------------------------
0 | l1   | Linear | 7.9 K
--------------------------------
7.9 K     Trainable params
0         Non-trainable params
7.9 K     Total params
0.031     Total estimated model params size (MB)
Epoch 1: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1875/1875 [00:07<00:00, 261.53it/s, loss=1.3, v_num=10]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

总结

以上我们用 30 行左右代码，实现了一个最简的 pytorch lightning 训练过程。这足以体现出 pytorch lightning 的简洁、易用。但是，显然这个最简实现缺少了很多东西，比如验证、测试、日志打印、模型保存等。接下来，我们将实现相对完整但依旧简洁的 pytorch lightning 模型开发过程。

pytorch lightning更多功能

本节将介绍相对更完整的 pytorch lightning 模型开发过程。

LighningModeul需实现方法

在一个相对完整的 LightnintModule 中，用户应当实现以下方法：

1 模型定义 (__init__)

通常定义模型的各个层，在 forward 调用这些层，完成模型前向。与原生 pytorch 类似。

2 前向计算 (forward)

与 torch.nn.Module 的 forward 中做的事情一样，调用 _init_ 中定义的层。完成模型前向。与原生 pytorch 类似。

3 训练/验证/测试步 (training_step/validation_step/test_step)

定义训练/测试/训练每一步中要做的事情，一般是计算损失、指标并返回。

def training_step(self, batch, batch_idx):
    # ....
    return xxx # 如果是training_step, 则必须包含损失
1
2
3

通常有两个入参 batch 和 batch_idx。是 batch 是 dataloader 给出的输入数据和标签，batch_idx 是当前 batch 的索引。

注意训练步的返回值必须是损失值，或者是包含 ‘loss’ 字段的字典。验证/测试步的返回值不必包括损失，可以是任意结果。

4 训练/验证/测试步结束后 (training_step_end/validation_step_end/test_step_end)

只在使用多个node进行训练且结果涉及如softmax之类需要全部输出联合运算的步骤时使用该函数。

5 训练/验证/测试轮结束后 (training_epoch_end/validation_epoch_end/test_epoch_end)

以 training_epoch_end 为例，其他类似。

如果需要对整一轮的结果进行处理，比如计算一些平均指标等，可以通过 training_epoch_end 来实现。

def training_epoch_end(self, outputs):
    # ....
    return xxx
1
2
3

其中入参 outputs 是一个列表，包含了每一步 training_step 返回的内容。我们可以在每一轮结束后，对每一步的结果进行处理。

4 选用优化器 (configure_optimizers)

设置模型参数更新所用的优化器。值得一提的是如果需要多个优化器（比如在训练 GAN 时），可以返回优化器列表。也可以在优化器的基础上返回学习率调整器，那就要返回两个列表。

5 数据加载器 (train_dataloader, val_dataloader, test_dataloader)

返回 dataloader。

各个 dataloader 也可以在运行 fit/validation/test 时传入，如：

train_loader = DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
model = MNISTModel()		# 不需要实现get_dataloader方法
trainer.fit(model, train_loader)
1
2
3

LightningModule自带工具

LightningModule 中提供了一些常用工具供用户直接使用：

log

Tensorboard 损失/指标日志保存和查看，不要自己定义，直接用即可。用法非常简单，将要记录的值传入：

self.log('train loss', loss)
1

当然一个功能完整的日志保存接口肯定提供了很多参数来控制，比如是按照 epoch 记录还是按照 step 记录、多卡训练时如何同步、指标是否要展示在进度条上、指标是否要保存在日志文件中等等。pytorch lightning 为这些选项都提供了控制参数，读者可以参考官方文档中 log 相关部分。

print

python 自带的 print 函数在进行多进程训练时会在每个进程都打印内容，这是原生 pytorch 进行分布式训练时一个很小但是很头疼的问题。LightningModule 提供的 print 只打印一次。

freeze

冻结所有权重以供预测时候使用。仅当已经训练完成且后面只测试时使用。

Trainer实例化参数

在实例化 Trainer 时，pytorch lightning 也提供了很多控制参数，这里介绍常用的几个，完整参数及含义请参考官方文档中 Trainer 相关部分。

• default_root_dir：默认存储地址。所有的实验变量和权重全部会被存到这个文件夹里面。默认情况下，实验结果会存在 lightning_logs/version_x/ 。
• max_epochs：最大训练周期数，默认为 1000，如果不设上限 epoch 数，设置为 -1。
• auto_scale_batch_size：在进行训练前自动选择合适的batch size。
• auto_select_gpus：自动选择合适的GPU。尤其是在有GPU处于独占模式时候，非常有用。
• gpus：控制使用的GPU数。当设定为None时，使用 cpu。
• auto_lr_find：自动找到合适的初始学习率。使用了该论文的技术。当且仅当执行 trainer.tune(model) 代码时工作。
• precision：浮点数精度。默认 32，即常规单精度 fp32 旬来呢。指定为 16 可以使用 fp16 精度加快模型训练并减少显存占用。
• val_check_interval：进行验证的周期。默认为 1，如果要训练 10 个 epoch 进行一次验证，设置为 10。
• fast_dev_run：如果设定为true，会只执行一个 batch 的 train, val 和 test，然后结束。仅用于debug。
• callbacks：需要调用的 callback 函数列表，关于常用 callback 函数下面会介绍。
• …

callback函数

Callback 是一个自包含的程序，可以与训练流程交织在一起，而不会污染主要的研究逻辑。Callback 并不一定只能在 epoch 结尾调用。pytorch-lightning 提供了数十个hook（接口，调用位置）可供选择，也可以自定义callback，实现任何想实现的模块。

推荐使用方式是，随问题和项目变化的操作，实现到 lightning module里面。而独立的、可复用的内容则可以定义单独的模块，方便多个模型调用。

常见的内建 callback 如：EarlyStopping，根据某个值，在数个epoch没有提升的情况下提前停止训练。。PrintTableMetricsCallback，在每个epoch结束后打印一份结果整理表格等。更多内建 callbacks 可参考相关文档。

模型加载与保存

模型保存

ModelCheckpoint 是一个自动储存的 callback 模块。默认情况下训练过程中只会自动储存最新的模型与相关参数，而用户可以通过这个 module 自定义。如观测一个 val_loss 的值，并储存 top 3 好的模型，且同时储存最后一个 epoch 的模型，等等。例：

from pytorch_lightning.callbacks import ModelCheckpoint

# saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckpt
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    filename='sample-mnist-{epoch:02d}-{val_loss:.2f}',
    save_top_k=3,
    mode='min',
    save_last=True
)

trainer = pl.Trainer(gpus=1, max_epochs=3, callbacks=[checkpoint_callback])
1
2
3
4
5
6
7
8
9
10
11
12

ModelCheckpoint Callback中，如果 save_weights_only=True，那么将会只储存模型的权重，相当于 model.save_weights(filepath)，反之会储存整个模型（包括模型结构），相当于model.save(filepath)）。

另外，也可以手动存储checkpoint: trainer.save_checkpoint("example.ckpt")

模型加载

加载一个模型，包括它的模型权重和超参数：

model = MyLightingModule.load_from_checkpoint(PATH)

print(model.learning_rate)
# 打印出超参数

model.eval()
y_hat = model(x)
1
2
3
4
5
6
7

加载模型时替换一些超参数：

class LitModel(LightningModule):
    def __init__(self, in_dim, out_dim):
      super().__init__()
      self.save_hyperparameters()
      self.l1 = nn.Linear(self.hparams.in_dim, self.hparams.out_dim)

# 如果在训练和保存模型时，超参数设置如下，在加载后可以替换这些超参数。
LitModel(in_dim=32, out_dim=10)

# 仍然使用in_dim=32, out_dim=10
model = LitModel.load_from_checkpoint(PATH)

# 替换为in_dim=128, out_dim=10
model = LitModel.load_from_checkpoint(PATH, in_dim=128, out_dim=10)
1
2
3
4
5
6
7
8
9
10
11
12
13
14

完整加载训练状态，包括模型的一切，以及和训练相关的一切参数，如 model, epoch, step, LR schedulers, apex 等。

model = LitModel()
trainer = Trainer(resume_from_checkpoint='some/path/to/my_checkpoint.ckpt')

# 自动恢复 model, epoch, step, LR schedulers, apex, etc...
trainer.fit(model)
1
2
3
4
5

实例

基于第三节介绍的更多功能，我们扩展第二节 MNIST 训练程序。代码如下。

import os

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
import pytorch_lightning as pl
import numpy as np


class MNISTModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.fc(x.view(-1, 28 * 28)))

    def training_step(self, batch, batch_nb):
        # REQUIRED
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log('train_loss', loss, on_step=False, on_epoch=True)
        return {'loss': loss}

    def validation_step(self, batch, batch_nb):
        # OPTIONAL
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        pred = y_hat.argmax(dim=1, keepdim=True)
        correct = pred.eq(y.view_as(pred)).sum().item()
        acc = correct / x.shape[0]
        self.log('val_acc', acc, on_step=False, on_epoch=True)
        self.log('val_loss', loss, on_step=False, on_epoch=True)
        return {'val_loss': loss, 'val_acc': acc}

    def validation_epoch_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_acc = np.mean([x['val_acc'] for x in outputs])
        return {'val_loss': avg_loss, 'val_acc': avg_acc}

    def test_step(self, batch, batch_nb):
        # OPTIONAL
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return {'test_loss': loss}

    def test_epoch_end(self, outputs):
        # OPTIONAL
        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
        return {'test_loss': avg_loss}

    def configure_optimizers(self):
        # REQUIRED
        return torch.optim.Adam(self.parameters(), lr=0.02)

    def train_dataloader(self):
        # REQUIRED
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

    def val_dataloader(self):
        # OPTIONAL
        return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)

    def test_dataloader(self):
        # OPTIONAL
        return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)

model = MNISTModel()
trainer = pl.Trainer(
        gpus=1,
        max_epochs=10,
        callbacks=[
            pl.callbacks.EarlyStopping( monitor="val_loss", patience=3),
        ]
)
trainer.fit(model)
trainer.test()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

Ref

• pytorch lightning 的官方文档
• Pytorch Lightning 完全攻略
• 参考代码