“简约版”Pytorch —— Pytorch-Lightning详解

介绍

来源：http://blog.itpub.net/31555081/viewspace-2698296/

PyTorch很容易使用，可以用来构建复杂的AI模型。但是一旦研究变得复杂，并且将诸如多GPU训练，16位精度和TPU训练之类的东西混在一起，用户很可能会写出有bug的代码。
PyTorch Lightning完全解决了这个问题。Lightning会构建您的PyTorch代码，以便抽象出训练的详细信息。这使得AI研究可扩展并且可以快速迭代。
PyTorch Lightning是NYU和FAIR(Facebook AI research)为从事AI研究的专业研究人员和博士生所创建的。

安装

conda activate your-env
pip install pytorch-lightning

实用功能 (Trainer参数详解)

来源：Lightning官方教程 和 官方文档

自动获取Batch Size - Automatic Batch Size Finder

auto_scale_batch_size

Batch Size一般会对模型的训练结果有影响，一般越大的batch size模型训练的结果会越好。有时候我们不知道自己的模型在当前的机器上最多能用多大的batch size，这时候通过Lightning Trainer的这个flag就可以帮助我们找到最大的batch size。

model = ...
# 设置为True，Trainer就会依次尝试用2的幂次方的batch size，直到超出内存
trainer = pl.Trainer(auto_scale_batch_size=True)
trainer.fit(model)

# 设置为'binsearch'，Trainer会用Binary Search的方式帮你找到最大的Batch Size
trainer = pl.Trainer(auto_scale_batch_size='binsearch')
trainer.tune(model)

# 注意：如果要用这个功能，在Module里面的__init__()函数中要有:
self.batch_size = batch_size
# 或者在__init__()里面调用:
self.save_hyperparameters()

自动获取初始学习率 - Automatic Learning Rate Finder

auto_lr_find

学习率learning rate是很重要的一个超参，选取一个合适的初始学习率也是很重要的，Lightning提供了这个有用的flag。
(15.11.2020)目前只支持单优化器(Optimizer)，预计在未来的几个月内会支持多优化器

import pytorch_lightning as pl

model = ...
# 可以直接设置为True，Trainer会自动用不同的学习率运行model，然后画出loss和学习率的曲线，帮你找到最合适的学习率
trainer = pl.Trainer(auto_lr_find=True)
trainer.tune(model)
print(model.learning_rate)

# 有时候我们会在model中给学习率起其他的名字，比如：
self.my_learning_rate = lr
# 这个时候我们可以用变量名直接设置auto_lr_find：
trainer = pl.Trainer(auto_lr_find='my_learing_rate')
# 开始寻找合适的学习率
lr_finder = trainer.tuner.lr_find(model)
# 展示loss和学习率的曲线
fig = lr_finder.plot(suggest=True)
fig.show()
# 设置为推荐的学习率
model.hyparams.learning_rate = lr_finder.suggestion()
# 开始训练
model.fit(model, train_loader, val_loader)

重新加载数据 - Reload DataLoaders Every Epoch

reload_dataloaders_every_epoch

一般数据只会在一开始加载一次，即在epochs的循环前面加载一次，然后每个循环都会shuffle之类的(如果你设置shuffle为True的话)。有时候我们的数据在训练过程中是会改变的，这个时候我们就需要在每个epoch都要再加载一次数据，Lightning就提供了这样一个flag，将其设置为True即可。

# 相当于：
# if False (default)
train_loader = model.train_dataloader()
for epoch in epochs:
    for batch in train_loader:
        ...

# if True
for epoch in epochs:
    train_loader = model.train_dataloader()
    for batch in train_loader:

回调函数 - Callbacks

callbacks

回调函数 (Callbacks) 在机器学习中也是很重要的工具，一般可以用来进行模型的断点断续，模型权重的存储，提早停止 (Early stop)，动态调整训练参数以及tensorboard之类的训练可视化等等。Lightning也支持非常灵活的Callbacks，只需要把Callbacks放进flag：callbacks中即可。Lightning提供了一些built-in的callbacks，同样也支持自定义callbacks类，所以非常灵活。
官方的callbacks说明

# 自定义Callback类
from pytorch_lightning.callbacks import Callback

class MyPrintingCallback(Callback):

    def on_init_start(self, trainer):
        print('Starting to init trainer!')

    def on_init_end(self, trainer):
        print('trainer is init now')

    def on_train_end(self, trainer, pl_module):
        print('do something when training ends')

trainer = Trainer(callbacks=[MyPrintingCallback()])

# 使用built-in的Callbacks
from pytorch_lightning.callbacks import EarlyStopping

# 可以直接使用默认的Callbacks
trainer = pl.Trainer(callbacks=[EarlyStopping('val_loss')])

# 也可以自己设置Callbacks的参数
early_stop = EarlyStopping(
	monitor='val_loss',
	patience=3,
	strict=False,
	verbose=False,
	mode='min'
)
trainer = pl.Trainer(callbacks=[early_stop])

展示网络信息 - Weights Summary

weights_summary

一般我们搭完模型都会去检查模型结构，计算模型的大小以及占用内存，之前需要安装一个额外的torchsummary的包可以用来输出pytorch模型的信息，Lightning则直接集成了这个功能，它提供三个参数：‘top’，'full’和None。

# (默认) 只展示顶层的模块
trainer = pl.Trainer(weights_summary='top')
# 展示所有的模块，包括子模块
trainer = pl.Trainer(weights_summary='full')
# 不打印模型信息
trainer = pl.Trainer(weights_summary=None)

进度条 - Progress Bar

progress_bar_refresh_rate

进度条 (Progress Bar) 是非常有用的工具，在python中我们一般使用第三方包tqdm来创建进度条，Lightning则集成了进度条，对应的值为steps，即每隔多少steps更新一次进度条，设置为0即不使用进度条。

# 默认每个step更新一次
trainer = Trainer(progress_bar_refresh_rate=1)

# 关闭进度条
trainer = Trainer(progress_bar_refresh_rate=0)

训练以及测试循环 - Training and Eval Loops

min_epochs, max_epochs, min_steps和max_steps

训练的时候我们需要指定Epoch的数量，Lightning提供了相应的flags。
对min_steps和max_steps而言，中间validation的steps不计入其中。
如果同时指定max_epochs和max_steps，则max_steps优先。

trainer = Trainer(min_steps=100, max_steps=10000)
trainer = Trainer(min_epochs=10, max_epochs=100)

check_val_every_n_epochs

训练过程中一般我们也会做测试 (validation) 来看模型的泛化能力以及是否过拟合。Lightning提供了对应的flag，默认为1，即每个epoch结束都会做一次validation。有时候一个training epoch只需要几分钟，这时候没有必要每个epoch结束都去做一次validation，我们可以设置为10或者任意其他的整数。

# 每5个epochs做一次validation
trainer = Trainer(check_val_every_n_epochs=5)

val_check_interval

对于某些大模型或者大数据集的训练，一个epoch可能会持续几个小时甚至几天，我们当然不能过这么久才做一次validation，这时候我们要用到这个新flag，它支持0到1的浮点数以及整数。比如0.25表示每隔四分之一的epoch做一次validation，1000表示每隔1000个steps做一次validation。

# 每隔四分之一个epoch做一次validation
trainer = Trainer(val_check_interval=0.25)
# 每隔1000个steps做一次validation
trainer = Trainer(val_check_interval=1000)

num_sanity_val_steps

一般validation都会在training之后做，如果validation的代码中存在bugs，而training一个epoch可能持续很久，这时候会浪费大量的时间。Lightning提供一个flag来解决这个问题，它会在正式training开始前先做若干个batches的validation，如果没有问题才会正式训练。它支持-1, 0以及正整数，-1表示做完整个validation，0表示不做。

# (默认) 正式训练前先做2个batches的validation
trainer = Trainer(num_sanity_val_steps=2)
# 做完一整个validation epoch才开始训练
trainer = Trainer(num_sanity_val_steps=-1)
# 直接开始训练
trainer = Trainer(num_sanity_val_steps=0)

limit_train_batches, limit_val_batches和limit_test_batches

有时候我们为了检查代码或者做测试必须跑完一整个或者更多的epochs，如果一个epoch的时间过长则会浪费时间，Lightning提供了解决方案，它们支持0到1的浮点数和整数，比如0.1代表每个epoch只跑十分之一的数据，10代表每个epoch只跑10个batches。

# 每个训练epoch只跑十分之一的数据
trainer = Trainer(limit_train_batches=0.1)
# 每个测试epoch只跑10个batches
trainer = Trainer(limit_val_batches=10)

单GPU以及多GPUs训练 - Training on GPUs

以前GPU基本用在游戏上，而现在被广泛应用在深度学习领域，因为GPU能高效地做矩阵运算，运算速度比CPU快进几百倍。而随着网络趋于复杂，计算量越来越大，一台GPU根本不够用，这时候就要多GPUs共同训练，Pytorch本身也支持多节点多GPUs训练，比如之前的DataParallel和DistributedDataParallel，但是开发人员需要自行修改很多代码，对刚刚接触分布式训练的开发人员不太友好，而Lightning则大大简化了这个过程，我们只需去掉代码中的**.cuda()和.to(device)，初始化Tensor的时候使用type_as和register_buffer**，然后在为Trainer添加对应的flag：gpus。

# 使用Lightning之前：
def forward(self, x):
	z = torch.Tensor(2, 3)
	z = z.cuda(0)

# 使用Lightning之后：
def forward(self, x):
	z = torch.Tensor(2, 3)
	z = z.type_as(x, device=self.device)

###########################################
# 使用一个GPU
trainer = pl.Trainer(gpus=1)
# 使用8个GPU
trainer = pl.Trainer(gpus=4)
# 使用所有可用的GPU
trainer = pl.Trainer(gpus=-1)
# 指定使用哪几个GPU
trainer = pl.Trainer(gpus=[0,3,4])

auto_select_gpus

有时候我们不知道哪些GPUs是没有被占用的，也就没办法指定GPU，Lightning为此提供了flag，它可以替我们检测可以使用的GPU个数以及序号。

# 不自动选择GPU (直接选择系统中的前两个GPU, 如果它们被占用则会失败)
trainer = Trainer(gpus=2, auto_select_gpus=False)

# 自动从系统中选择两个可用的GPU
trainer = Trainer(gpus=2, auto_select_gpus=True)

# 指定所有的GPU，不管它们是否被占用
Trainer(gpus=-1, auto_select_gpus=False)

# 指定所有可用的GPU (如果只有一个GPU可用，则只会使用一个GPU)
Trainer(gpus=-1, auto_select_gpus=True)

log_gpu_memory

如果我们想要监测GPU内存的使用状况，Lightning也提供了相应的flag。使用的话可能会使训练变慢，因为它使用的是nvidia-smi的输出。

# 默认不监测GPU内存
trainer = Trainer(log_gpu_memory=None)

# 监测主节点上的所有GPU
trainer = Trainer(log_gpu_memory='all')

# 只记录主节点上的最小以及最大GPU内存使用
trainer = Trainer(log_gpu_memory='min_max')

benchmark

如果你模型的输入大小保持不变，那么可以设置cudnn.benchmark为True，这样可以加速训练，如果输入大小不固定，那么反而会减慢训练。

# 默认为False
trainer = Trainer(benchmark=False)

deterministic

如果你想保证模型的再现性 (reproducibility)，那么应该设置cudnn.deterministic为True，虽然会减慢速度。

# 默认为False
trainer = Trainer(deterministic=False)

进阶分布式训练 - Advanced distributed training

replace_sampler_ddp

进行多GPU训练时需要为DataLoader指定Sample，来为GPU分配数据。Lightning默认会进行此操作，对训练集shuffle，测试集不shuffle，如果你想改变sampler，可以设置replace_sampler_ddp为False，然后自己指定Sampler。

# 默认为True
trainer = Trainer(replace_sampler_ddp=True)

# 可以自己指定Sampler
trainer = Trainer(replace_sampler_ddp=False)
sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=True)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

prepare_data_per_node

如果多节点训练即多机器训练，默认会在每个节点上调用prepare_data()，如果这些节点公用一个文件系统，那么可以设为False，这样只会在主节点的主GPU上调用prepare_data()。

# 默认为True
Trainer(prepare_data_per_node=True)

# use only NODE_RANK=0, LOCAL_RANK=0
Trainer(prepare_data_per_node=False)

sync_batchnorm

Batch Normalization对模型的训练影响很大，所以默认会同步所有GPU上的Batch Normalization，这对Batch Size比较小的模型训练更重要。

trainer = Trainer(sync_batchnorm=True)

代码测试 - Debugging

fast_dev_run

如果你想在正式训练前检查下所有代码，那么设置fast_dev_run为True，这样Trainer会为训练验证测试分别跑一个batch，从而找到潜在的Bugs。它不会产生logs和checkpoints

# 默认为False
trainer = Trainer(fast_dev_run=False)

# 为训练验证测试分别跑一个batch，然后程序结束
trainer = Trainer(fast_dev_run=True)

overfit_batches

有时候为了快速测试代码或者想要model达到过拟合 (测试模型有效性)，可以为其设置一个数，可以是0到1的浮点数或者整数。

# 默认为0.0
trainer = Trainer(overfit_batches=0.0)

# 只用1%的训练集训练以及验证和测试
trainer = Trainer(overfit_batches=0.01)

# 在10个batches上过拟合
trainer = Trainer(overfit_batches=10)

梯度累积 - Accumulating Gradients

accumulate_grad_batches

特定情况下，数据太大没法放进一个batch里，或者需要大的batch size而没有那么多GPU可以用，这时可以通过累计梯度的方式模拟大的batch size。

# 默认不累积梯度
trainer = Trainer(accumulate_grad_batches=1)

# 每次累积四个batches的梯度 (batch size相当于增大四倍)
trainer = Trainer(accumulate_grad_batches=4)

# epochs 1-4 不累积梯度，epochs 5-10 累积3个batches的梯度，之后累积20个batches的梯度
trainer = Trainer(accumulate_grad_batches={5: 3, 10: 20})

混合精度训练 - Mixed Precision Training

现在的模型精度默认都是32-bit，论文MIXED PRECISION TRAINING指出使用混合精度训练 (32-bit + 16-bit) 在保证精度的情况下减少了内存的占用，可以使用更大的batch size，训练也会更快。

precision

可以指定训练精度，默认32.

amp_level

使用NVIDIA的apex框架，优化级别有00, 01, 02, 03，其中00代表全精度，03代表半精度，具体可以参考Apex官方文档

# 默认全精度 - 32-bit
trainer = Trainer(precision=32)
# 半精度 - 16-bit
trainer = Trainer(precision=16)

# 默认使用02，即'Almost FP16'
trainer = Trainer(amp_backend='apex', amp_level='O2')

示例代码

来源：https://github.com/PyTorchLightning/deep-learning-project-template

from argparse import ArgumentParser

import torch
import pytorch_lightning as pl
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split

from torchvision.datasets.mnist import MNIST
from torchvision import transforms


class LitClassifier(pl.LightningModule):
    def __init__(self, hidden_dim=128, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()

        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))
        x = torch.relu(self.l2(x))
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log('valid_loss', loss)

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        self.log('test_loss', loss)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument('--hidden_dim', type=int, default=128)
        parser.add_argument('--learning_rate', type=float, default=0.0001)
        return parser


def cli_main():
    pl.seed_everything(1234)

    # ------------
    # args
    # ------------
    parser = ArgumentParser()
    parser.add_argument('--batch_size', default=32, type=int)
    parser = pl.Trainer.add_argparse_args(parser)
    parser = LitClassifier.add_model_specific_args(parser)
    args = parser.parse_args()

    # ------------
    # data
    # ------------
    dataset = MNIST('', train=True, download=True, transform=transforms.ToTensor())
    mnist_test = MNIST('', train=False, download=True, transform=transforms.ToTensor())
    mnist_train, mnist_val = random_split(dataset, [55000, 5000])

    train_loader = DataLoader(mnist_train, batch_size=args.batch_size)
    val_loader = DataLoader(mnist_val, batch_size=args.batch_size)
    test_loader = DataLoader(mnist_test, batch_size=args.batch_size)

    # ------------
    # model
    # ------------
    model = LitClassifier(args.hidden_dim, args.learning_rate)

    # ------------
    # training
    # ------------
    trainer = pl.Trainer.from_argparse_args(args)
    trainer.fit(model, train_loader, val_loader)

    # ------------
    # testing
    # ------------
    trainer.test(test_dataloaders=test_loader)


if __name__ == '__main__':
    cli_main()