Pytorch深度学习模型训练加速
本文主要解读在Pytorch进行神经网络深度学习时,导致训练时间大大增加的一些应当被摒弃的坏习惯,以及一些能够加速训练的方法。从而让我们在使用Pytorch的时候更加快速。
这个指南的原作者LORENZ KUHN是来自ETH计算机科学的研究生,在此鸣谢。
此外,我参考了AI公园的内容,非商业用途,若有侵权请告知我删除。
他给出的方案中,会专注于可以直接在PyTorch中进行的更改,而不需要引入额外的库,但是我们假设要使用GPU训练模型。
1. 考虑使用另外一种学习率策略
学习率对收敛速度和模型的泛化性能有很大影响,因此采用循环学习率和1Cycle学习率策略(proposed by Leslie N. Smith)。其中,1Cycle学习率策略变化曲线是这样的:
这种策略可以实现巨大的加速 —— Smith称之为“超级收敛”。在ImageNet的ResNet-56训练迭代数上,使用1Cycle策略减少了1/10。
PyTorch也是很贴心的实现了这两个方法:
torch.optim.lr_scheduler.CyclicLR
torch.optim.lr_scheduler.OneCycleLR
这两个策略的一个缺点是它们引入了许多额外的超参数。
为什么这样做可以提升收敛速度呢?一个可能的解释是定期提高学习率有助于更快的穿越“鞍点”。
2. 在 DataLoader中使用多个workers和pinned memory
当使用torch.utils.data.DataLoader
时:
设置num_workers
> 0,设置pin_memory=True
但是增加num_workers
会增加CPU内存消耗,因此需要注意CPU内存和参数的协同配置。
3. 最大化batch size(博主个人不推荐)
这是一个颇有争议的观点。一般来说,使用GPU允许的最大的batch size会加速训练。
但是如果修改batch大小,还必须调整其他超参数,例如学习率。
这里的一个经验法则是:
当你把batch数量翻倍时,学习率也要翻倍。
但是,使用大batch可能会导致泛化能力变差。
4. 使用自动混合精度(还没弄懂,待更)
5. 考虑使用另外的优化器:AdamW
AdamW是具有权重衰减的Adam。
可在PyTorch中直接使用:torch.optim.AdamW
。
无论在误差还是训练时间上,AdamW都比Adam表现更好。
此外,还有一些自带优化器最近受到了很多关注,最著名的是LARS和LAMB。
6. 开启cudNN benchmarking
若模型架构保持不变,可以设置torch.backends.cudnn.benchmark = True
需要注意,如果像第3点中的方法将batch size
最大化,那么这种自动调优可能会变得非常缓慢。
7. 注意CPU和GPU之间频繁的数据传输
如果创建新的张量,可使用关键字参数device=torch.device('cuda:0')
直接将它分配给GPU。
如果确实需要传输数据,在传输后使用.to(non_blocking=True)
可能会有用。
8. 使用gradient/activation检查点(待更)
9. 使用梯度累加
可以在调用optimizer.step()之前,在多个.backward()中累积梯度。
梯度累加可以实现如下(by 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
10. 对于多个GPU使用分布式数据并行
对于分布式训练加速,一个简单的方法是使用torch.nn.DistributedDataParallel
而不是torch.nn.DataParallel
。通过这样做,每个GPU将由一个专用的CPU核心驱动,避免了DataParallel的GIL问题。
11. 将梯度设为None而不是0
这一点改起来很方便,效果也很好,但是可能会有伴随的副作用!!!
使用.zero_grad(set_to_none=True)
而不是.zero_grad()
。这样做会让内存分配器去处理梯度,而不是主动将它们设置为0。正如在文档中所说的那样,这会导致产生一个适度的加速,所以不要期待任何奇迹。
12. 使用.as_tensor() 而不是 .tensor()
torch.tensor()
会拷贝数据,如果一个numpy数组想转为tensor,使用 torch.as_tensor()
或 torch.from_numpy()
来避免拷贝数据。
13. 需要的时候打开调试工具
Pytorch提供了大量的有用的调试工具,如autograd.profiler
,autograd.grad_check
和autograd.anomaly_detection
。在需要的时候使用它们,在不需要它们的时候关闭它们,因为它们会减慢你的训练。
14. 使用梯度剪裁(待更)
15. 在BatchNorm之前不使用bias
这个方法效果很一般,但这是一个非常简单的方法:在BatchNormalization
层之前不使用bias
。对于二维卷积层,可以将关键字bias设为False:
torch.nn.Conv2d(..., bias=False, ...)
16. 在验证的时候关闭梯度计算
这一点很多人在调试之后忘记了这一点,导致GPU根本无法发挥作用。
这个很直接:在验证的时候使用 torch.no_grad()
。
17. 对输入和batch使用归一化
请检查一下是否在使用batch-normalization