【阅读笔记】针对PyTorch提高神经网络训练速度的方法—《Here are 17 ways of making PyTorch training faster – what did I miss?》

原文连接：https://efficientdl.com/faster-deep-learning-in-pytorch-a-guide/

写在开头

在这篇文章中，我将概述一些加速Pytorch深度学习模型训练的最低投入，最大的影响方式。对于每种方法，我将简要总结这个想法，尝试估算预期的加速并讨论一些局限性。我将专注于传达最重要的部分，并指向每个部分的进一步资源。大多数情况下，我将重点关注可以直接在Pytorch内进行的更改，而无需引入其他库，我会假设您正在训练GPU上的模型。
以下建议按照提高速度效果从高到低排序

1. 考虑使用不同的学习率变化策略（Consider using a different learning rate schedule）
2. 在DataLoader中使用多线程并且固定内存（Use multiple workers and pinned memory in
   DataLoader）
3. BatchSize最大化（Max out the batch size）
4. 使用自动混合精度（Use Automatic Mixed Precision (AMP)）
5. 使用不同的优化器（Consider using a different optimizer）
6. 打开cudNN基准测试（Turn on cudNN benchmarking）
7. 防止CPU和GPU之间频繁的传输数据（Beware of frequently transferring data between
   CPUs and GPUs）
8. 使用梯度/激活检查点（Use gradient/activation checkpointing）
9. 使用梯度累积（Use gradient accumulation）
10. 使用DistributedDataParallel进行多GPU训练（Use DistributedDataParallel for
    multi-GPU training）
11. 将梯度设置为None而不是0（Set gradients to None rather than 0）
12. 使用 .as_tensor 而不是 .tensor()（Use .as_tensor rather than .tensor()）
13. 如果不需要，请关闭调试API（Turn off debugging APIs if not needed）
14. 使用梯度剪裁（Use gradient clipping）
15. 在BatchNorm之前关闭偏置（Turn off bias before BatchNorm）
16. 在验证过程中关闭梯度计算（Turn off gradient computation during validation）
17. 使用输入和批次归一化（Use input and batch normalization）

此外还有笔者自己平时在实验中的一些建议
18. 数据变换 (用于数据增强) 可成为速度提升的另一个来源。一些只使用简单 Python 语句的变换可以通过使用 numba 包来加速。
19. 注意排查速度慢的瓶颈是什么。

1. 考虑使用不同的学习率变化策略（Consider using a different learning rate schedule）

选择的学习率时间表对收敛速度以及模型的泛化性能有很大影响。
Leslie Smith提出的周期性学习速率(CLR)以及 1cycle 策略可以令复杂模型的训练迅速完成。
论文连接：
Cyclical Learning Rates for Training Neural Networks
Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates
1cycle策略的学习率表如下

[1循环包括]两个相等长度的步骤，一个从较低的学习率到更高的步骤，而不是回到最小值。最大值应该是通过学习率查找器选择的值，较低的值可能会降低十倍。然后，该周期的长度应略低于epochs的总数，在训练的最后一部分中，我们应允许学习率降低到最低限度的数量级。
在最好的情况下，与传统的相比，这个策略实现了大规模的提速。不过有一个缺点，它们引入了一些额外的超参数。
为什么这样做有效？一种可能的解释是，定期增加学习率有助于更快地穿越损失函数中的鞍点。
PyTorch提供了两种方法的实现：
torch.optim.lr_scheduler.CyclicLR、torch.optim.lr_scheduler.OneCycleLR
这些调度程序的一个缺点是他们引入了许多其他超参数。下面的文章和仓库提供了一个很好的概述和实施，以了解如何找到良好的超参数，包括上述学习率查找器。
文章：Hyper-parameter Tuning Techniques in Deep Learning（需要开VPN）
仓库：https://github.com/davidtvs/pytorch-lr-finder

2. 在DataLoader中使用多线程并且固定内存（Use multiple workers and pinned memory in DataLoader）

使用时torch.utils.data.DataLoader，请设置num_workers > 0，而不是默认值0，和pin_memory=True，而不是默认值False。
英伟达高级工程师Szymon Micacz使用了4个工作程序和固定内存，在单个训练时期内将速度提高了两倍。
需要注意的是，在选择worker数量时，建议将设置为可用GPU数量的四倍。worker数量的多和少都会导致速度变慢，数量越多还会增加CPU内存消耗。

3. BatchSize最大化（Max out the batch size）

这是一个有争议的点。但是，通常，似乎使用最大的批量大小，您的GPU内存许可将加速您的培训（例如，请参阅Nvidia的Szymon Migacz）。请注意，如果修改批处理大小，您还必须调整其他超参数，例如学习率。一项经验法则是将学习率翻倍，使学习率加倍。
但是，使用大批量尺寸的缺点之一是，它们可能导致解决方案比受过较小批次训练的解决方案更糟。

4. 使用自动混合精度（Use Automatic Mixed Precision (AMP)）

PyTorch 1.6版本就包括了对 PyTorch 的自动混合精度训练的本地实现。
与其他地方使用的单精度（FP32）相比，某些操作可以在半精度（FP16）上运行得更快，并且不会损失准确性。
随后，让AMP自动决定应以什么样的格式执行操作，这样既可以加快训练速度，也可以减少内存占用。
在最好的情况下，AMP的使用如下：

import torch
# Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler()

for data, label in data_iter:
   optimizer.zero_grad()
   # Casts operations to mixed precision
   with torch.cuda.amp.autocast():
      loss = model(data)

   # Scales the loss, and calls backward()
   # to create scaled gradients
   scaler.scale(loss).backward()

   # Unscales gradients and calls
   # or skips optimizer.step()
   scaler.step(optimizer)

   # Updates the scale for next iteration
   scaler.update()

代码来自：https://pytorch.org/docs/stable/amp.html#gradient-scaling
有研究者发现，在NVIDIA V100 GPU上对一些常见的语言和视觉模型进行基准测试时，使用AMP要比常规的FP32训练的速度提升2倍，最高可提升5.5倍。
目前，只有CUDA ops 可以通过这种方式进行自动广播。具体的限制和细节见如下文档
文档：https://pytorch.org/docs/stable/amp.html#op-eligibility

5. 使用不同的优化器（Consider using a different optimizer）

比如AdamW，AdamW是带有权重衰减（而不是L2正则化）的Adam，它在错误实现、训练时间都胜过Adam。
AdamW，Adam都可以结合1Cycle policy使用。
此外，还有一些非本地的优化器值得关注，比如，LARS和LAMB。
NVIDA的APEX实现了一些常见优化器（比如Adam）的融合版本，比如Adam。与Adam的PyTorch实现相比，它避免了多次进出GPU内存的过程，产生了5%左右的速度提升。

6. 打开cudNN基准测试（Turn on cudNN benchmarking）

如果你的模型架构保持固定，输入大小保持不变，则可以设置torch.backends.cudnn.benchmark = True，启动 cudNN 自动调整器。
它将对cudNN中计算卷积的多种不同方法进行基准测试，以获得最佳的性能指标。
这里有一个警告是，如果您最大化上述批处理大小，则这种自动调用可能会变得非常慢。

7.防止CPU和GPU之间频繁的传输数据（Beware of frequently transferring data between CPUs and GPUs）

请谨慎使用tensor.cpu()和tensor.cuda()，这些将消耗很多的时间，同样的.item()和.numpy()也是一样。这些可以使用.detach()代替。
如果正在创建一个张量，就可以使用关键字参数device=torch.device(‘cuda:0’)直接将其分配给你的GPU。

8. 使用梯度/激活检查点（Use gradient/activation checkpointing）

检查点的工作原理，是用计算换取内存。检查点部分不是讲整个计算图的所有中间激活都存储起来向后计算，而不是保存中间激活，在后传中重新计算。
它可以应用到模型的任何部分。
具体来说，在前向传递中，函数将以torch.no_grad()的方式运行，即不存储中间的激活。相反，前向传递会保存输入元组和函数参数。
在后向传递中，检索保存的输入和函数，然后再次对函数进行前向传递计算，现在跟踪中间激活，使用这些激活值计算梯度。
虽然这可能会略微增加你在给定批量大小下的运行时间，但你会显著减少你的内存占用。这反过来又会让你进一步增加你所使用的批次大小，提高GPU的利用率。
虽然检查点本地用作torch.utils.checkpoint，但似乎确实需要一些思考和精力才能正常实施。Priya Goyal有一个很好的教程，展示了检查点的一些关键方面。
教程：https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb

9. 使用梯度累积（Use gradient accumulation）

另一种增加批次大小的方法是在调用optimizer.step()之前，在多个.backward()通道中累积梯度。
具体实施过程如下：

model.zero_grad()                                   # Reset gradients tensors
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # Forward pass
    loss = loss_function(predictions, labels)       # Compute loss function
    loss = loss / accumulation_steps                # Normalize our loss (if averaged)
    loss.backward()                                 # Backward pass
    if (i+1) % accumulation_steps == 0:             # Wait for several backward steps
        optimizer.step()                            # Now we can do an optimizer step
        model.zero_grad()                           # Reset gradients tensors
        if (i+1) % evaluation_steps == 0:           # Evaluate the model when we...
            evaluate_model()                        # ...have no gradients accumulated

参考：https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
这个方法主要是为了规避GPU内存限制而开发的，但不清楚是否有额外的.backward()循环之间的权衡。

10. 使用DistributedDataParallel进行多GPU训练（Use DistributedDataParallel for multi-GPU training）

这个方法比较常见，如果具体来看的话可能需要单独写一篇文章，但一个简单的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。
这样做可以让每个GPU将由一个专门的CPU核驱动，避免了DataParallel的GIL问题。

11. 将梯度设置为None而不是0（Set gradients to None rather than 0）

使用.zero_grad(set_to_none=True)而不是.zero_grad()。
这样做会让内存分配器来处理梯度，而不是主动将它们设置为0，这样会适度加速。
注意，这样做并不是没有副作用的。细节请见文档

12. 使用 .as_tensor 而不是 .tensor()（Use .as_tensor rather than .tensor()）

torch.tensor() 总是复制数据。如果你有一个要转换的 numpy 数组，使用 torch.as_tensor() 或 torch.from_numpy() 来避免复制数据。

13. 如果不需要，请关闭调试API（Turn off debugging APIs if not needed）

Pytorch提供了很多调试工具，例如autograd.profiler，autograd.grad_check和autograd.anomaly_detection，确保在需要的时候使用它们，不需要时将其关闭，否则他们会拖慢你的训练速度。

14. 使用梯度剪裁（Use gradient clipping）

剪裁梯度，可以加速加速收敛。最初是用来避免RNNs中的梯度爆炸，可以使用torch.nn.utils.clipgrad_norm来实现。
目前尚不清楚哪些模型能靠梯度剪裁能够加速多少，但它似乎对RNNs、基于 Transformer 和 ResNets 的架构以及一系列不同的优化器都非常有用。
huggingface的transformer代码中有很好的使用举例。

15. 在BatchNorm之前关闭偏置（Turn off bias before BatchNorm）

这是一个非常简单的方法，在BatchNormalization图层之前关闭图层的偏置。
对于二维卷积层，可以通过将bias关键字设置为False：来完成torch.nn.Conv2d(…, bias=False, …)

16. 在验证过程中关闭梯度计算（Turn off gradient computation during validation）

在验证期间设置torch.no_grad() 。

17. 使用输入和批次归一化（Use input and batch normalization）

您可能已经这样做了，但是您可能需要仔细检查：

• 是否对输入进行了归一化操作
• 是否已经使用了batch normalization

18. 数据变换 (用于数据增强) 可成为速度提升的另一个来源。一些只使用简单 Python 语句的变换可以通过使用 numba 包来加速。

19. 注意排查速度慢的瓶颈是什么

除了GPU和CPU等硬件的问题导致训练速度慢之外，也要注意排查磁盘IO的问题。可以输出网络的数据加载时间和网络训练时间来判断。
此外，频繁的向LOG文件中写入日志也将导致训练时间打折扣，建议间隔多个iter写入一次，比如50个iter写入一次log文件。