pytorch_lightning模型训练 | Civilization Museum

type

status

password

date

slug

summary

背景

pytorch-lightning 是建立在pytorch之上的高层次模型接口。pytorch-lightning 之于 pytorch，就如同keras之于 tensorflow.pytorch-lightning 有以下一些引人注目的功能：

可以不必编写自定义循环，只要指定loss计算方法即可。

可以通过callbacks非常方便地添加CheckPoint参数保存、early_stopping 等功能。

可以非常方便地在单CPU、多CPU、单GPU、多GPU乃至多TPU上训练模型。

可以通过调用torchmetrics库，非常方便地添加Accuracy,AUC,Precision等各种常用评估指标。

可以非常方便地实施多批次梯度累加、半精度混合精度训练、最大batch_size自动搜索等技巧，加快训练过程。

可以非常方便地使用SWA(随机参数平均)、CyclicLR(学习率周期性调度策略)与auto_lr_find(最优学习率发现)等技巧实现模型涨点。

pl_docs: https://pytorch-lightning.readthedocs.io/en/latest/starter/introduction.html
pl_template: https://github.com/PyTorchLightning/deep-learning-project-template
torchmetrics: https://torchmetrics.readthedocs.io/en/latest/pages/lightning.html

一、pytorch-lightning的设计哲学

pytorch-lightning 的核心设计哲学是将深度学习项目中的研究代码(定义模型) 和工程代码 (训练模型) 相互分离。用户只需专注于研究代码(pl.LightningModule)的实现，而工程代码借助训练工具类(pl.Trainer)统一实现。更详细地说，深度学习项目代码可以分成如下4部分：

研究代码 (Research code)，用户继承LightningModule实现。

工程代码 (Engineering code)，用户无需关注通过调用Trainer实现。

非必要代码（Non-essential research code，logging, etc...），用户通过调用Callbacks实现。

数据 (Data)，用户通过torch.utils.data.DataLoader实现，也可以封装成pl.LightningDataModule。

二、pytorch-lightning使用范例

下面我们使用minist图片分类问题为例，演示pytorch-lightning的最佳实践。

1、准备数据

示例化

2、定义模型

tensor.detach()：返回一个新的tensor，从当前计算图中分离下来的，但是仍指向原变量的存放位置,不同之处只是requires_grad为false，得到的这个tensor永远不需要计算其梯度，不具有grad。

tensor.detach_()：将一个tensor从创建它的图中分离，并把它设置成叶子 tensor 其实就相当于变量之间的关系本来是 x -> m -> y，这里的叶子 tensor 是 x，但是这个时候对 m 进行了m.detach_()操作,其实就是进行了两个操作：

将 m 的 grad_fn 的值设置为 None,这样 m 就不会再与前一个节点 x 关联，这里的关系就会变成 x, m -> y,此时的 m 就变成了叶子结点
然后会将 m 的 requires_grad 设置为 False，这样对 y 进行 backward() 时就不会求 m的梯度

输出

💡

如果需要对每个training_step的输出做一些操作，可以通过改training_epoch_end来实现

3、训练模型

输出

4、评估模型

输出

5、使用模型

6、保存模型

最优模型默认保存在 trainer.checkpoint_callback.best_model_path 的目录下，可以直接加载。

输出

三、训练加速技巧

下面重点介绍pytorch_lightning 模型训练加速的一些技巧。

1，使用多进程读取数据(num_workers=4)

2，使用锁业内存(pin_memory=True)

3，使用加速器(gpus=4,strategy="ddp_find_unused_parameters_false")

4，使用梯度累加(accumulate_grad_batches=6)

5，使用半精度(precision=16,batch_size=2*batch_size)

6，自动搜索最大batch_size(auto_scale_batch_size='binsearch')

（注：过大的batch_size对模型学习是有害的。）详细原理，可以参考：https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html我们将训练代码封装成如下脚本形式，方便后面测试使用。

mnist_cnn.py

数据集

1、使用多进程读取数据(num_workers=4)

使用多进程读取数据，可以避免数据加载过程成为性能瓶颈。

单进程读取数据(num_workers=0, gpus=1): 1min 18s

多进程读取数据(num_workers=4, gpus=1): 59.7s

2、使用锁业内存(pin_memory=True)

锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘）因此锁业内存比非锁业内存读写效率更高，copy到GPU上也更快速。当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

非锁业内存存储数据(pin_memory=False, gpus=1): 1min

锁业内存存储数据(pin_memory=True, gpus=1): 59.5s

输出

训练

1、使用加速器(gpus=4,strategy="ddp_find_unused_parameters_false")

pl 可以很方便地应用单CPU、多CPU、单GPU、多GPU乃至多TPU上训练模型。以下几种情况训练耗时统计如下：

单CPU: 2min 17s

单GPU: 59.4 s

4个GPU(dp模式): 1min

4个GPU(ddp模式): 38.9 s

一般情况下，如果是单机多卡，建议使用 ddp模式，因为dp模式需要非常多的data和model传输，非常耗时。

参数	备注
`gpus=0`	`单CPU模式`
`num_processes=4,strategy="ddp_find_unused_parameters_false"`	`多CPU(进程)模式`
`gpus=1`	`单GPU模式`
`gpus=4,strategy="dp"`	`多GPU(dp速度提升效果一般)`
`gpus=4,strategy=“ddp_find_unused_parameters_false"`	`多GPU(ddp速度提升效果好)`

2、使用梯度累加(accumulate_grad_batches=6)

梯度累加就是累加多个batch的梯度，然后用累加的梯度更新一次参数，使用梯度累加相当于增大batch_size.由于更新参数的计算量略大于简单梯度求和的计算量（对于大部分优化器而言），使用梯度累加会让速度略有提升。

4个GPU(ddp模式): 38.9 s

4个GPU(ddp模式)+梯度累加: 36.9 s

3、使用半精度(precision=16)

通过precision可以设置 double (64), float (32), bfloat16 ("bf16"), half (16) 精度的训练。默认是 float(32) 标准精度，bfloat16 ("bf16") 是混合精度。如果选择 half(16) 半精度，并同时增大 batch_size 为原来2倍, 通常训练速度会提升3倍左右。

4、自动搜索最大batch_size(auto_scale_batch_size="power")

四，训练涨分技巧

pytorch_lightning 可以非常容易地支持以下训练涨分技巧:

SWA(随机参数平均): pl.callbacks.stochastic_weight_avg.StochasticWeightAveraging

CyclicLR(学习率周期性调度策略): 设置 lr_scheduler 为 torch.optim.lr_scheduler.CyclicLR实现。

auto_lr_find最优学习率发现: 设置 pl.Trainer(auto_lr_find = True) 实现。

💡

参考论文:

Cyclical Learning Rates for Training Neural Networks.pdf

965.1KB

Averaging Weights Leads to Wider Optima and Better Generalization.pdf

716.0KB

我们将代码整理成如下形式，以便后续测试使用。