『PyTorch』PyTorch深度学习模型训练加速指南

---

简要介绍在 PyTorch 中加速深度学习模型训练的一些最小改动、影响最大的方法。我既喜欢效率又喜欢ML，所以我尝试进行记录。

1. 考虑使用另外一种学习率策略

你选择的学习率对收敛速度以及模型的泛化性能有很大的影响。
循环学习率 和 1Cycle学习率策略都是 Leslie N. Smith 提出的方法，然后由fast.ai推广。

PyTorch实现了这两个方法，torch.optim.lr_scheduler.CyclicLR 和 torch.optim.lr_scheduler.OneCycleLR。

Note： 没有尝试使用过，待后期尝试后更新。

2. 在 DataLoader 中使用多个 workers 和 pinned memory

当使用 torch.utils.data.DataLoader 时，设置 num_workers > 0，而不是等于0，设置 pin_memory=True 而不是默认值False。详细解释：https://pytorch.org/docs/stable/data.html。

Szymon Micacz通过使用 4 个 workers 和 pinned memory，实现了单个训练 epoch 的 2 倍加速。

一个经验法则，选择 workers 的数量设置为可用GPU数量的 4 倍，更大或更小的 workers 数量会变慢。

Note： 增加 num_workers 会增加 CPU 内存消耗。

3. 最大化 batch size

这是一个颇有争议的观点。一般来说，似乎使用 GPU 允许的最大的 batch size 可能会加速你的训练。

Note： 如果你修改了 batch size，你还必须调整其他超参数，例如学习率。这里的一个经验法则是，当你把 batch 数量翻倍时，学习率也要翻倍。

OpenAI有一篇很好的实证论文关于不同 batch size 需要的收敛步骤的数量。Daniel Huynh运行一些实验用不同 batch size(使用上面所讨论的1Cycle策略)，从 batch size 64 到 512 他实现了 4 倍的加速。

然而，使用大 batch 的缺点之一是，它们可能会导致泛化能力比使用小 batch 的模型差。

4. 使用自动混合精度

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()

5. 开启 cudnn benchmarking

如果你的模型架构保持不变，你的输入大小保持不变，设置 torch.backends.cudnn.benchmark = True 可能是有益的。这使得 cudNN 能够测试许多不同的卷积计算方法，然后使用最快的方法。

对于加速的预期有一个粗略的参考，Szymon Migacz 达到 70% 的 forward 的加速以及 27% 的 forward 和 backward 的加速。

这里需要注意的是，如果你像上面提到的那样将 batch size 最大化，那么这种自动调优可能会变得非常缓慢。

6. 注意 CPU 和 GPU 之间频繁的数据传输

小心使用 tensor.cpu() 和 tensor.cuda() 频繁地将张量从 GPU 和 CPU 之间相互转换。对于 .item() 和 .numpy() 也是一样，用 .detach() 代替。

如果你正在创建一个新的张量，你也可以使用关键字参数 device=torch.device('cuda:0') 直接将它分配给你的 GPU。

如果你确实需要传输数据，在传输后使用 .to(non_blocking=True) 可能会很有用，只要你没有任何同步点。

如果你真的需要，你可以试试 Santosh Gupta 的 SpeedTorch，虽然不是很确定在什么情况下可以加速。

7. 使用 gradient/activation 检查点

直接引用文档中的话：

检查点的工作原理是用计算交换内存，并不是存储整个计算图的所有中间激活用于向后计算，检查点不保存中间的激活，而是在向后传递中重新计算它们。可以应用于模型的任何部分。
具体来说，在向前传递中，function 会以 torch.no_grad() 的方式运行，也就是说，不存储中间激活。相反，正向传递保存输入和 function 的参数。在向后传递中，将检索保存的输入和 function，并再次根据 function 计算向前传递，然后跟踪中间的激活，再使用这些激活值计算梯度。

因此，虽然这可能会略微增加给定 batch 大小的运行时间，但会显著减少内存占用。这反过来会允许你进一步增加你正在使用的 batch 大小，从而更好地利用 GPU。

检查点的 pytorch 实现为 torch.utils.checkpoint，需要想点办法才能实现的很好。

8. 使用梯度累加

增加 batch 大小的另一种方法是在调用 optimizer.step() 之前，在多个 .backward() 中累积梯度。

在Hugging Face的实现中，梯度累加可以实现如下：

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

这个方法主要是为了避开 GPU 内存限制。fastai论坛上的这个讨论：https://forums.fast.ai/t/accumulating-gradients/33219/28 似乎表明它实际上可以加速训练，所以可能值得一试。

9. 对于多个GPU使用分布式数据并行

对于分布式训练加速，一个简单的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。通过这样做，每个 GPU 将由一个专用的 CPU 核心驱动，避免了 DataParallel 的 GIL 问题。

10. 将梯度设为 None 而不是 0

使用 .zero_grad(set_to_none=True) 而不是 .zero_grad()。这样做会让内存分配器去处理梯度，而不是主动将它们设置为0。正如在文档中所说的那样，这会导致产生一个适度的加速，所以不要期待任何奇迹。

Note： 这样做并不是没有副作用的！关于这一点的详细信息请查看文档。

11. 使用.as_tensor() 而不是 .tensor()

torch.tensor() 会拷贝数据，如果你有一个 numpy 数组，你想转为 tensor，使用 torch.as_tensor() 或是 torch.from_numpy() 来避免拷贝数据。

12. 使用梯度剪裁

最初是用于 RNNs 避免爆炸梯度，有一些经验证据和一些理论支持认为剪裁梯度(粗略地说： gradient = min(gradient, threshold)) 可以加速收敛。Hugging Face 的 Transformer 实现是关于如何使用梯度剪裁以及其他的一些方法如 AMP 的一个非常干净的例子。

在 PyTorch 中，这可以通过使用 torch.nn.utils.clip_grad_norm_ 实现。我并不完全清楚哪个模型从梯度裁剪中获益多少，但它似乎对 RNN、基于 Transformer 和 ResNets 架构以及一系列不同的优化器都非常有用。

13. 在 BatchNorm 之前不使用 bias

这是一个非常简单的方法：在 BatchNormalization 层之前不使用 bias。对于二维卷积层，可以将关键字 bias 设为 False: torch.nn.Conv2d(..., bias=False, ...)。

你会保存一些参数，然而，与这里提到的其他一些方法相比，我对这个方法的加速期望相对较小。

14. 在验证的时候关闭梯度计算

这个很直接：在验证的时候使用 torch.no_grad()。

15. 对输入和 batch 使用归一化

参考文献

1. 关于学习率策略的：https://sgugger.github.io/the-1cycle-policy.html

2. 关于找最佳学习率的：https://sgugger.github.io/how-do-you-find-a-good-learning-rate.html

3. AdamW相关的：https://www.fast.ai/2018/07/02/adam-weight-decay/

4. https://mp.weixin.qq.com/s/OHGvIW-oqh2ijtEeg9Seng