近日，Reddit 上一个帖子热度爆表。主题内容是关于怎样加速 PyTorch 训练。原文作者是来自苏黎世联邦理工学院的计算机科学硕士生 LORENZ KUHN，文章向我们介绍了在使用 PyTorch 训练深度模型时最省力、最有效的 17 种方法。 该文所提方法，都是假设你在 GPU 环境下训练模型。具体内容如下。 17 种加速 PyTorch 训练的方法 1. 考虑换一种学习率 schedule 学习率 schedule 的选择对模型的收敛速度和泛化能力有很大的影响。Leslie N. Smith 等人在论文《Cyclical Learning Rates for Training Neural Networks》、《Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates 》中提出了周期性（Cyclical）学习率以及 1Cycle 学习率 schedule。之后，fast.ai 的 Jeremy Howard 和 Sylvain Gugger 对其进行了推广。下图是 1Cycle 学习率 schedule 的图示： Sylvain 写到：1Cycle 包括两个等长的步幅，一个步幅是从较低的学习率到较高的学习率，另一个是回到最低水平。最大值来自学习率查找器选取的值，较小的值可以低十倍。然后，这个周期的长度应该略小于总的 epochs 数，并且，在训练的最后阶段，我们应该允许学习率比最小值小几个数量级。 与传统的学习率 schedule 相比，在最好的情况下，该 schedule 实现了巨大的加速（Smith 称之为超级收敛）。例如，使用 1Cycle 策略在 ImageNet 数据集上训练 ResNet-56，训练迭代次数减少为原来的 1/10，但模型性能仍能比肩原论文中的水平。在常见的体系架构和优化器中，这种 schedule 似乎表现得很好。 Pytorch 已经实现了这两种方法：「torch.optim.lr_scheduler.CyclicLR」和「torch.optim.lr_scheduler.OneCycleLR」。 参考文档：https://pytorch.org/docs/stable/optim.html 2. 在 DataLoader 中使用多个 worker 和页锁定内存 当使用 torch.utils.data.DataLoader 时，设置 num_workers > 0，而不是默认值 0，同时设置 pin_memory=True，而不是默认值 False。 参考文档：https://pytorch.org/docs/stable/data.html 来自 NVIDIA 的高级 CUDA 深度学习算法软件工程师 Szymon Micacz 就曾使用四个 worker 和页锁定内存（pinned memory）在单个 epoch 中实现了 2 倍的加速。人们选择 worker 数量的经验法则是将其设置为可用 GPU 数量的四倍，大于或小于这个数都会降低训练速度。请注意，增加 num_workers 将增加 CPU 内存消耗。 3. 把 batch 调到最大 把 batch 调到最大是一个颇有争议的观点。一般来说，如果在 GPU 内存允许的范围内将 batch 调到最大，你的训练速度会更快。但是，你也必须调整其他超参数，比如学习率。一个比较好用的经验是，batch 大小加倍时，学习率也要加倍。 OpenAI 的论文《An Empirical Model of Large-Batch Training》很好地论证了不同的 batch 大小需要多少步才能收敛。在《How to get 4x speedup and better generalization using the right batch size》一文中，作者 Daniel Huynh 使用不同的 batch 大小进行了一些实验（也使用上面讨论的 1Cycle 策略）。最终，他将 batch 大小由 64 增加到 512，实现了 4 倍的加速。 然而，使用大 batch 的不足是，这可能导致解决方案的泛化能力比使用小 batch 的差。 4. 使用自动混合精度（AMP） PyTorch 1.6 版本包括对 PyTorch 的自动混合精度训练的本地实现。这里想说的是，与单精度 (FP32) 相比，某些运算在半精度 (FP16) 下运行更快，而不会损失准确率。AMP 会自动决定应该以哪种精度执行哪种运算。这样既可以加快训练速度，又可以减少内存占用。 在最好的情况下，AMP 的使用情况如下：
importtorch# Creates once at the beginning of trainingscaler = torch.cuda.amp.GradScaler() fordata, label in data_iter:optimizer.zero_grad() # Casts operations to mixed precisionwithtorch.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 iterationscaler.update()
model.zero_grad() # Reset gradients tensorsfor 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 passif (i+1) % accumulation_steps == 0: # Wait for several backward steps optimizer.step() # Now we can do an optimizer step model.zero_grad() # Reset gradients tensorsif (i+1) % evaluation_steps == 0: # Evaluate the model when we... evaluate_model() # ...have no gradients accumulate