全文共 4857 字, 预计学习时长 10 分钟
事实上, 你的模型可能还停留在石器时代的水平. 估计你还在用 32 位精度或 * GASP(一般活动仿真语言)* 训练, 甚至可能只在单 GPU 上训练. 如果市面上有 99 个加速指南, 但你可能只看过 1 个?(没错, 就是这样). 但这份终极指南, 会一步步教你清除模型中所有的(GP 模型).
这份指南的介绍从简单到复杂, 一直介绍到你可以完成的大多数 PITA 修改, 以充分利用你的网络. 例子中会包括一些 Pytorch 代码和相关标记, 可以在 Pytorch-Lightning 训练器中用, 以防大家不想自己敲码!
这份指南针对的是谁? 任何用 Pytorch 研究非琐碎的深度学习模型的人, 比如工业研究人员, 博士生, 学者等等...... 这些模型可能要花费几天, 甚至几周, 几个月的时间来训练.
指南(从易到难)
1. 使用 DataLoader.
2. DataLoader 中的进程数.
3. 批尺寸.
4. 累积梯度.
5. 保留计算图.
6. 转至单 GPU.
7. 16 位混合精度训练.
8. 转至多 GPU(模型复制).
9. 转至多 GPU 节点(8+GPUs).
10. 有关模型加速的思考和技巧
Pytorch-Lightning
文中讨论的各种优化, 都可以在名为 Pytorch-Lightning(https://github.com/williamFalcon/pytorch-lightning?source=post_page---------------------------) 的 Pytorch 图书馆中找到.
Lightning 是基于 Pytorch 的一个光包装器, 它可以帮助研究人员自动训练模型, 但关键的模型部件还是由研究人员完全控制.
参照此篇教程, 获得更有力的范例(https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/single_gpu_node_template.py?source=post_page---------------------------).
Lightning 采用最新, 最尖端的方法, 将犯错的可能性降到最低.
MNIST 定义的 Lightning 模型(https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py?source=post_page---------------------------), 可适用于训练器.
- from pytorch-lightning import Trainer
- model = LightningModule(...)
- trainer = Trainer()
- trainer.fit(model)
1. DataLoader
这可能是最容易提速的地方. 靠保存 h5py 或 numpy 文件来加速数据加载的日子已经一去不复返了. 用 Pytorch dataloader (https://pytorch.org/tutorials/beginner/data_loading_tutorial.html?source=post_page---------------------------)加载图像数据非常简单.(关于 NLP 数据, 请参照 TorchText:https://torchtext.readthedocs.io/en/latest/datasets.html?source=post_page---------------------------)
- dataset = MNIST(root=self.hparams.data_root, train=train, download=True)
- loader = DataLoader(dataset, batch_size=32, shuffle=True)
- for batch in loader:
- x, y = batch
- model.training_step(x, y)
- ...
在 Lightning 中, 你无需指定一个训练循环, 只需定义 dataLoaders, 训练器便会在 需要时调用它们(https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py?source=post_page---------------------------#L163-L217).
2. DataLoaders 中的进程数
加快速度的第二个秘诀在于允许批量并行加载. 所以, 你可以一次加载许多批量, 而不是一次加载一个.
- # slow
- loader = DataLoader(dataset, batch_size=32, shuffle=True)
- # fast (use 10 workers)
- loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)
3. 批量大小(Batch size)
在开始下一步优化步骤之前, 将批量大小调高到 CPU 内存或 GPU 内存允许的最大值.
接下来的部分将着重于减少内存占用, 这样就可以继续增加批尺寸.
记住, 你很可能需要再次更新学习率. 如果将批尺寸增加一倍, 最好将学习速度也提高一倍.
4. 累积梯度
假如已经最大限度地使用了计算资源, 而批尺寸仍然太低(假设为 8), 那我们则需为梯度下降模拟更大的批尺寸, 以供精准估计.
假设想让批尺寸达到 128. 然后, 在执行单个优化器步骤前, 将执行 16 次前向和后向传播(批量大小为 8).
- # clear last step
- optimizer.zero_grad()
- # 16 accumulated gradient steps
- scaled_loss = 0
- for accumulated_step_i in range(16):
- out = model.forward()
- loss = some_loss(out,y)
- loss.backward()
- scaled_loss += loss.item()
- # update weights after 8 steps. effective batch = 8*16
- optimizer.step()
- # loss is now scaled up by the number of accumulated batches
- actual_loss = scaled_loss / 16
而在 Lightning 中, 这些已经自动执行了. 只需设置标记: https://williamfalcon.github.io/pytorch-lightning/Trainer/Training Loop/?source=post_page---------------------------#accumulated-gradients
- trainer = Trainer(accumulate_grad_batches=16)
- trainer.fit(model)
5. 保留计算图
撑爆内存很简单, 只要不释放指向计算图形的指针, 比如...... 为记录日志保存 loss.
- losses = []
- ...
- losses.append(loss)
- print(f'current loss: {torch.mean(losses)'})
上述的问题在于, loss 仍然有一个图形副本. 在这种情况中, 可用. item()来释放它.
- # bad
- losses.append(loss)
- # good
- losses.append(loss.item())
Lightning 会特别注意, 让其无法保留图形副本 (示例: https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/trainer.py?source=post_page---------------------------#L767-L768)
6. 单 GPU 训练
一旦完成了前面的步骤, 就可以进入 GPU 训练了. GPU 的训练将对许多 GPU 核心上的数学计算进行并行处理. 能加速多少取决于使用的 GPU 类型. 个人使用的话, 推荐使用 2080Ti, 公司使用的话可用 V100.
刚开始你可能会觉得压力很大, 但其实只需做两件事: 1)将你的模型移动到 GPU 上, 2)在用其运行数据时, 把数据导至 GPU 中.
- # put model on GPU
- model.cuda(0)
- # put data on gpu (cuda on a variable returns a cuda copy)
- x = x.cuda(0)
- # runs on GPU now
- model(x)
如果使用 Lightning, 则不需要对代码做任何操作. 只需设置标记(https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed training/?source=post_page---------------------------#single-gpu):
- # ask lightning to use gpu 0 for training
- trainer = Trainer(gpus=[0])
- trainer.fit(model)
在 GPU 进行训练时, 要注意限制 CPU 和 GPU 之间的传输量.
- # expensive
- x = x.cuda(0)
- # very expensive
- x = x.CPU()
- x = x.cuda(0)
例如, 如果耗尽了内存, 不要为了省内存, 将数据移回 CPU. 尝试用其他方式优化代码, 或者在用这种方法之前先跨 GPUs 分配代码.
此外还要注意进行强制 GPUs 同步的操作. 例如清除内存缓存.
- # really bad idea.Stops all the GPUs until they all catch up
- torch.cuda.empty_cache()
但是如果使用 Lightning, 那么只有在定义 Lightning 模块时可能会出现这种问题. Lightning 特别注意避免此类错误.
7. 16 位精度
16 位精度可以有效地削减一半的内存占用. 大多数模型都是用 32 位精度数进行训练的. 然而最近的研究发现, 使用 16 位精度, 模型也可以很好地工作. 混合精度指的是, 用 16 位训练一些特定的模型, 而权值类的用 32 位训练.
要想在 Pytorch 中用 16 位精度, 先从 NVIDIA 中安装 apex 图书馆 并对你的模型进行这些更改.
- # enable 16-bit on the model and the optimizer
- model, optimizers = amp.initialize(model, optimizers, opt_level='O2')
- # when doing .backward, let amp do it so it can scale the loss
- with amp.scale_loss(loss, optimizer) as scaled_loss:
- scaled_loss.backward()
amp 包会处理大部分事情. 如果梯度爆炸或趋于零, 它甚至会扩大 loss.
在 Lightning 中, 使用 16 位很简单(https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed training/?source=post_page---------------------------#16-bit-mixed-precision), 不需对你的模型做任何修改, 也不用完成上述操作.
- trainer = Trainer(amp_level='O2', use_amp=False)
- trainer.fit(model)
8. 移至多 GPU
现在, 事情就变得有意思了. 有 3 种 (也许更多?) 方式训练多 GPU.
分批量训练
A)在每个 GPU 上复制模型; B)给每个 GPU 分配一部分批量.
第一种方法叫做分批量训练. 这一策略将模型复制到每个 GPU 上, 而每个 GPU 会分到该批量的一部分.
- # copy model on each GPU and give a fourth of the batch to each
- model = DataParallel(model, devices=[0, 1, 2 ,3])
- # out has 4 outputs (one for each gpu)
- out = model(x.cuda(0))
在 Lightning 中, 可以直接指示训练器增加 GPU 数量, 而无需完成上述任何操作.
- # ask lightning to use 4 GPUs for training
- trainer = Trainer(gpus=[0, 1, 2, 3])
- trainer.fit(model)
分模型训练
将模型的不同部分分配给不同的 GPU, 按顺序分配批量
有时模型可能太大, 内存不足以支撑. 比如, 带有编码器和解码器的 Sequence to Sequence 模型在生成输出时可能会占用 20gb 的内存. 在这种情况下, 我们希望把编码器和解码器放在单独的 GPU 上.
- # each model is sooo big we can't fit both in memory
- encoder_rnn.cuda(0)
- decoder_rnn.cuda(1)
- # run input through encoder on GPU 0
- out = encoder_rnn(x.cuda(0))
- # run output through decoder on the next GPU
- out = decoder_rnn(x.cuda(1))
- # normally we want to bring all outputs back to GPU 0
- out = out.cuda(0)
对于这种类型的训练, 无需将 Lightning 训练器分到任何 GPU 上. 与之相反, 只要把自己的模块导入正确的 GPU 的 Lightning 模块中:
- class MyModule(LightningModule):
- def __init__():
- self.encoder = RNN(...)
- self.decoder = RNN(...)
- def forward(x):
- # models won't be moved after the first forward because
- # they are already on the correct GPUs
- self.encoder.cuda(0)
- self.decoder.cuda(1)
- out = self.encoder(x)
- out = self.decoder(out.cuda(1))
- # don't pass GPUs to trainer
- model = MyModule()
- trainer = Trainer()
- trainer.fit(model)
混合两种训练方法
在上面的例子中, 编码器和解码器仍然可以从并行化每个操作中获益. 我们现在可以更具创造力了.
- # change these lines
- self.encoder = RNN(...)
- self.decoder = RNN(...)
- # to these
- # now each RNN is based on a different gpu set
- self.encoder = DataParallel(self.encoder, devices=[0, 1, 2, 3])
- self.decoder = DataParallel(self.encoder, devices=[4, 5, 6, 7])
- # in forward...
- out = self.encoder(x.cuda(0))
- # notice inputs on first gpu in device
- sout = self.decoder(out.cuda(4)) # <--- the 4 here
使用多 GPUs 时需注意的事项
. 如果该设备上已存在 model.cuda(), 那么它不会完成任何操作.
. 始终输入到设备列表中的第一个设备上.
. 跨设备传输数据非常昂贵, 不到万不得已不要这样做.
. 优化器和梯度将存储在 GPU 0 上. 因此, GPU 0 使用的内存很可能比其他处理器大得多.
9. 多节点 GPU 训练
每台机器上的各 GPU 都可获取一份模型的副本. 每台机器分得一部分数据, 并仅针对该部分数据进行训练. 各机器彼此同步梯度.
做到了这一步, 就可以在几分钟内训练 Imagenet 数据集了! 这没有想象中那么难, 但需要更多有关计算集群的知识. 这些指令假定你正在集群上使用 SLURM.
Pytorch 在各个 GPU 上跨节点复制模型并同步梯度, 从而实现多节点训练. 因此, 每个模型都是在各 GPU 上独立初始化的, 本质上是在数据的一个分区上独立训练的, 只是它们都接收来自所有模型的梯度更新.
高级阶段:
1. 在各 GPU 上初始化一个模型的副本(确保设置好种子, 使每个模型初始化到相同的权值, 否则操作会失效.)
2. 将数据集分成子集. 每个 GPU 只在自己的子集上训练.
3. On .backward() 所有副本都会接收各模型梯度的副本. 只有此时, 模型之间才会相互通信.
Pytorch 有一个很好的抽象概念, 叫做分布式数据并行处理, 它可以为你完成这一操作. 要使用 DDP(分布式数据并行处理), 需要做 4 件事 :
- def tng_dataloader():
- d = MNIST()
- # 4: Add distributed sampler
- # sampler sends a portion of tng data to each machine
- dist_sampler = DistributedSampler(dataset)
- dataloader = DataLoader(d, shuffle=False, sampler=dist_sampler)
- def main_process_entrypoint(gpu_nb):
- # 2: set up connections between all gpus across all machines
- # all gpus connect to a single GPU "root"
- # the default uses env://
- world = nb_gpus * nb_nodes
- dist.init_process_group("nccl", rank=gpu_nb, world_size=world)
- # 3: wrap model in DPP
- torch.cuda.set_device(gpu_nb)
- model.cuda(gpu_nb)
- model = DistributedDataParallel(model, device_ids=[gpu_nb])
- # train your model now...
- if __name__ == '__main__':
- # 1: spawn number of processes
- # your cluster will call main for each machine
- mp.spawn(main_process_entrypoint, nprocs=8)
Pytorch 团队对此有一份详细的实用教程(https://github.com/pytorch/examples/blob/master/imagenet/main.py?source=post_page---------------------------).
然而, 在 Lightning 中, 这是一个自带功能. 只需设定节点数标志, 其余的交给 Lightning 处理就好.
- # train on 1024 gpus across 128 nodes
- trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])
Lightning 还附带了一个 SlurmCluster 管理器, 可助你简单地提交 SLURM 任务的正确细节(示例: https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/multi_node_cluster_template.py?source=post_page---------------------------#L103-L134)
10. 福利! 更快的多 GPU 单节点训练
事实证明, 分布式数据并行处理要比数据并行快得多, 因为其唯一的通信是梯度同步. 因此, 最好用分布式数据并行处理替换数据并行, 即使只是在做单机训练.
在 Lightning 中, 通过将 distributed_backend 设置为 ddp(分布式数据并行处理)并设置 GPU 的数量, 这可以很容易实现.
- # train on 4 gpus on the same machine MUCH faster than DataParallel
- trainer = Trainer(distributed_backend='ddp', gpus=[0, 1, 2, 3])
有关模型加速的思考和技巧
如何通过寻找瓶颈来思考问题? 可以把模型分成几个部分:
首先, 确保数据加载中没有瓶颈. 为此, 可以使用上述的现有数据加载方案, 但是如果没有适合你的方案, 你可以把离线处理及超高速缓存作为高性能数据储存, 就像 h5py 一样.
接下来看看在训练过程中该怎么做. 确保快速转发, 避免多余的计算, 并将 CPU 和 GPU 之间的数据传输最小化. 最后, 避免降低 GPU 的速度(在本指南中有介绍).
接下来, 最大化批尺寸, 通常来说, GPU 的内存大小会限制批量大小. 自此看来, 这其实就是跨 GPU 分布, 但要最小化延迟, 有效使用大批次(例如在数据集中, 可能会在多个 GPUs 上获得 8000 + 的有效批量大小).
但是需要小心处理大批次. 根据具体问题查阅文献, 学习一下别人是如何处理的!
留言 点赞 发个朋友圈
我们一起分享 AI 学习与发展的干货
编译组: 曹娜, 夏伊凡
相关链接:
https://towardsdatascience.com/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565
如需转载, 请后台留言, 遵守转载规范
来源: http://www.tuicool.com/articles/ziyQzeQ