避免让你的神经网络变得如此低效
现实是，你的模型可能依然停留在较旧的技术水平。我敢肯定你仍在使用32位精度，甚至只在单个GPU上进行训练。
虽然网上有许多关于神经网络加速的指南，但通常缺乏一个完整的检查清单（现在有了），通过这个清单，你可以逐步确保充分利用模型的性能。
本指南涵盖从简单到复杂的各种优化策略，以帮助你的网络获得最佳效果。我将提供示例PyTorch代码以及可在PyTorch Lightning TRAIneR中使用的相关flags，这样你无需自己去编写这些代码！
**本指南适合谁？**任何使用PyTorch进行深度学习模型研究的人，包括研究人员、博士生、学者等，我们讨论的模型可能需要几天、几周甚至几个月的训练时间。
我们将讨论：
- 使用Dataloader
- Dataloader中的workers数量
- 批处理大小
- 梯度累加
- 保留的计算图
- 单个GPU训练
- 16位混合精度训练
- 多GPU训练
- 多节点GPU训练
- 思考模型加速的方法
PyTorch-Lightning
你可以在PyTorch库PyTorch-Lightning中找到我所讨论的每一个优化。Lightning是一个基于PyTorch的封装，可以自动化训练，并让研究人员完全控制关键的模型组件。它使用最新的最佳实践，最大限度地减少可能出错的地方。
我们为MNIST定义了LightningModel，并使用TRAIneR来训练模型。
1. Dataloader
这是获取速度提升的最简单方式。保存h5py或numpy文件以加速数据加载的时代已经过去，使用PyTorch Dataloader加载图像数据变得十分简单（对于NLP数据，请查看TorchText）。
在Lightning中，你不需要手动指定训练循环，只需定义Dataloader和TRAIneR即可在需要时调用它们。
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) …
2. Dataloader中的workers数量
另一个加速的关键在于允许批量并行加载。因此，你可以一次加载多个batch，而不是逐个加载。
3. 批处理大小
在进行下一个优化步骤之前，将批处理大小增大到CPU-RAM或GPU-RAM所允许的最大值。
下一节将重点介绍如何帮助减少内存占用，以便你能够继续增加批处理大小。
请记得你可能需要再次调整学习率。好的经验法则是，如果批处理大小加倍，学习率也应加倍。
4. 梯度累加
当你的批处理大小已经达到计算资源的上限时，比如只有8，你需要模拟一个更大的批处理大小来进行梯度下降，以获得良好的估计。
在Lightning中，这一切都已经为你准备好了，只需设置accumulate_grad_batches=16：
trainer = Trainer(accumulate_grad_batches=16) trainer.fit(Model)
5. 保留的计算图
一个常见的导致内存溢出的问题是为了记录日志而保存你的损失。
losses = [] … losses.append(loss) print(f’Current loss: {torch.mean(losses)}’)
问题在于，损失仍然包含整个计算图的副本。在这种情况下，可以调用.item()来释放它。
Lightning会非常小心，确保不会保留计算图的副本。
6. 单个GPU训练
完成前面的步骤后，是时候进行GPU训练了。在GPU上进行训练可以并行化多个GPU核心之间的数学计算。你获得的加速取决于所使用的GPU类型。我个人推荐使用2080Ti，公司使用V100。
乍一看，这可能会让你感到困惑，但实际上你只需做两件事：1）将模型移动到GPU，2）在运行数据时将数据放到GPU上。
如果使用Lightning，你无需做任何额外的事情，只需设置trainer(gpus=1)。
在GPU上训练时，主要需要注意的是限制CPU和GPU之间的数据传输次数。
如果内存耗尽，避免将数据移回CPU以节省内存。在求助于GPU之前，尝试其他方法优化代码或在GPU之间分配内存。
另一个注意事项是调用强制GPU同步的操作，例如清除内存缓存。
然而，如果使用Lightning，可能出现问题的地方在于定义Lightning模块时。Lightning会特别注意避免这种错误。
7. 16位精度
16位精度是一种显著减少内存占用的技术。大多数模型使用32位精度进行训练，但最近的研究发现，16位模型同样表现良好。混合精度意味着对某些部分使用16位，而权重等则保持在32位。
要在PyTorch中使用16位精度，请安装NVIDIA的apex库，并对模型进行相应的调整。
AMP包会处理大部分事情。如果出现梯度爆炸或趋近于0，它甚至会缩放损失。
在Lightning中，启用16位精度不需要修改模型中的任何内容，只需设置trainer(precision=16)即可。
8. 移动到多个GPU
此时，事情变得非常有趣。有三种（也许更多？）方法来进行多GPU训练。
- 分批训练
- 模型分布训练
- 两者的结合
9. 多节点GPU训练
每个机器上的每个GPU都有一个模型的副本。每台机器获得数据的一部分，并仅在该部分上进行训练。每台机器都能同步梯度。
如果你已经完成这一步，那么你现在可以在几分钟内训练ImageNet！这并不像你想象的那么复杂，但可能需要你对计算集群有更多了解。这些说明假设你在集群上使用SLURM。
PyTorch允许多节点训练，通过在每个节点上复制每个GPU上的模型并同步梯度。因此，每个模型都是在每个GPU上独立初始化的，基本上独立地在数据的一个分区上进行训练，除了它们都接收来自所有模型的梯度更新。
从高层次看：
- 在每个GPU上初始化一个模型的副本（确保设置种子，以使每个模型初始化为相同的权重，否则会失败）。
- 将数据集分割成子集（使用DistributedSampler）。每个GPU只在自己的小子集上进行训练。
- 在.backward()中，所有副本都接收所有模型的梯度副本。这是模型之间唯一的通信。
PyTorch有一个很好的抽象，称为DistributedDataParallel，可以帮助你实现这一功能。要使用DDP，你需要做四件事情：
然而，在Lightning中，只需设置节点数量，就会为你处理剩下的事情。
10. 福利！在单个节点上更快的多GPU训练
事实证明，DistributedDataParallel比DataParallel快得多，因为它只执行梯度同步的通信。因此，一个好的技巧是在单机训练时使用DistributedDataParallel替换DataParallel。
对模型加速的思考
虽然本指南为你提供了一系列提升网络速度的技巧，但我还是想强调如何通过查找瓶颈来思考问题。
我将模型分为几个部分：
首先，我确保数据加载中没有瓶颈。为此，我使用了前述的数据加载解决方案，但如果没有解决方案满足你的需求，请考虑离线处理并缓存到高性能数据存储中，比如h5py。
接下来，检查训练步骤中要完成的任务。确保前向传播速度快，避免过多计算，尽量减少CPU和GPU之间的数据传输。最后，避免做一些会降低GPU速度的操作（本指南中已有提及）。
接下来，我会尽量最大化批处理大小，这通常受限于GPU内存的大小。此时，需要关注如何在多个GPU上分配大的批处理大小，并尽量减少延迟（例如，我可能会尝试在多个GPU上使用8000+的有效批处理大小）。
然而，你需要谨慎处理大的批处理大小。针对具体问题，请查阅相关文献，看看人们可能忽视了哪些！
