我是靠谱客的博主 火星上冬天,最近开发中收集的这篇文章主要介绍加速Pytorch训练,觉得挺不错的,现在分享给大家,希望可以做个参考。

概述

转载自:Here are 17 ways of making PyTorch training faster-what did I miss?

按加速程度从大到小的大致排序为:

  1. 使用不同学习率的schedule;
  2. 在DataLoader中:使用多个worker&使用pin memory;
  3. 最大化batch size;
  4. 使用自动混合精度(AMP);
  5. 使用不同的优化器optimizer;
  6. 打开cudnn benchmark;
  7. 避免CPU和GPU之间频繁传输数据;
  8. 使用梯度/激活checkpointing;
  9. 使用累计梯度;
  10. 使用DistributedDataParallel进行多GPU训练;
  11. 设置梯度为None(而不是0);
  12. 使用 .as_tensor()(而不是.tensor());
  13. 关掉不必要的debugging API;
  14. 使用gradient clipping;
  15. BatchNorm前设置:bias=False;
  16. 在验证阶段关闭梯度计算;
  17. 输入归一化&批量标准化;

1 使用不同学习率的schedule

学习率会影响模型的收敛速度&&泛化性能。这里推荐两种:
(1)Cyclical learning rate
torch.optim.lr_scheduler.CyclicLR()
参考:Cyclical Learning Rates for Training Neural Networks

(2)1Cycle learning rate
torch.optim.lr_scheduler.OneCycleLR()
其schedule如下图:
1Cycle learning rate schedule

  • 1Cycle包含两个等长的步骤:①从一个较小的学习率增加到一个较大的学习率(1-175 epoch);②从较大的学习率回到小的学习率(175-350 epoch);
  • 学习率的较大值是由Learning Rate Finder选取;
  • 较大值 = 10*较小值;
  • 两个等长步骤加起来的迭代次数略小于总的epoch数,且在最后几个epoch中,学习率可比较小值还小几个数量级。

参考:Super-Convergence:Very Fast Training of Neural Networks Using Large Learning Rates

缺点

上述scheduler的缺点是:引入了额外的超参数。关于这些超参数的选取可以参考:Hyper-parameter Tuning Techniques in Deep Learning和Pytorch LR finder。

至于为什么上述scheduler可以加速训练,有一种可能的解释是:定期提高学习率有助于更快地遍历loss的鞍点。

2 在DataLoader中:使用多个worker&使用pin memory

当使用torch.utils.data.DataLoader时,设置num_workers > 0以及pin_memory=True

num_workers的选取

经验法则:可用GPU数量的4倍。
过多或过少的worker都会导致速度变慢。

缺点

增加worker的数量会增加CPU内存消耗(worker加载数据到RAM的进程是CPU复制)。

3 最大化batch size

这是一个比较有争议的做法。参考:An Empirical Model of Large-Batch Training和How to get 4× speedup and better generalization using the right batch size

注意

当修改batch size时,需要对其他超参数(譬如学习率)进行调整。
经验法则:当batch size翻倍时,学习率也翻倍。

缺点

降低泛化性能。

4 使用自动混合精度(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()

关于其作用原理可以参考Tensor Layouts In Memory:NCHW vs NHWC。

注意

使用AMP有某些限制,例如其现在仅限于CUDA操作,具体参考:Autocast Op Reference。

5 使用不同的优化器optimizer

torch.optim.AdamW代替torch.optim.Adam:前者使用weight decay而非L2正则化,其在误差和训练时间方面都优于后者。
参考:AdamW and Super-convergence is now the fastest way to train neural nets

其他的非本地优化器(not-yet-native optimizer)还有:LARSLAMB

6 打开cudnn benchmark

如果模型架构是固定的&&输入大小不变,则可以设置torch.backends.cudnn.benchmark = True
其启用了 cudNN autotuner,对 cudNN 中计算卷积的多种不同方法进行基准测试,然后使用最快的方法。
参考:pytorch docs: torch.backends.cudnn

注意

如果batch size过大,则这种自动调整可能会变得很慢。

7 避免CPU和GPU之间频繁传输数据

  • 避免频繁地使用 tensor.cpu() 将张量从 GPU 转到 CPU或者使用tensor.cuda() 将张量从 CPU 转到 GPU;
  • detach() 代替.item().numpy()
  • 创建新张量时,可以使用device=torch.device('cuda:0') 将其直接分配给 GPU;
  • 若需要传输数据,而且传输之后没有同步点,则可以使用.to(non_blocking=True)

8 使用梯度/激活checkpointing

checkpointing牺牲计算,换取内存:不使用checkpointing时,计算机会存储整个计算图中的所有中间激活用于反向计算;而checkpointed part不保存中间激活,而是在反向传播中重新计算它们。checkpointing可以应用于模型的任何部分。
在前向传播中,checkpointing将以torch.no_gard()方式运行,即不存储中间节点,而是保存输入元组和函数参数。在反向传播中,检索保存的输入和函数,再次在函数上计算前向传播,此时追踪中间激活,然后再使用这些激活值计算梯度。
综上,checkpointing可能会稍微增加运行时间,但是因为减少了内存占用,batch size有了增大空间,从而提高GPU的利用率。

torch.utils.checkpoint可以实现checkpointing,但是正确使用它是不容易,参考:checkpointing for PyTorch models

9 使用累计梯度

其可以视作一种增加batch size的方式:在调用 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 accumulate
注意

该方法主要是为了避免GPU内存受限,但它实际中可以加速训练。

10 使用DistributedDataParallel进行多GPU训练

加速分布式训练的方法:用torch.nn.DistributedDataParallel替代torch.nn.DataParallel,使得每个GPU由专用的CPU内核驱动,避免了后者的GIL问题。
参考:PYTORCH DISTRIBUTED OVERVIEW

GIL(global interpreter lock)问题

早期cpython(python最主要的解释器)为了实现多线程的功能,提高CPU利用率,“无脑”使用了GIL机制(全局锁)来解决线程之间数据同步问题(历史问题埋下的坑)。
在GIL的机制下,正常的CPU切换步骤中,在线程获得cpu时间片后,还需要获得GIL才能够执行。
在单核时代,这种机制是行得通的:因为只有在GIL释放之后才会触发OS的线程调度,那么这时候其它线程势必能够获取GIL,实现线程的切换执行。
但在多核情况下:因为只有一把大锁GIL,在多核架构下的线程调度,就会出现:虽然线程获得了cpu时间片,然而却没有得到GIL,只好白白浪费CPU时间片,然后等待切换进入等待,等待下一次被唤醒,如此恶性循环。那么在任何时候,只能有一个线程在执行,也就是说只有一个CPU核心真正执行代码。

11 设置梯度为None(而不是0)

.zero_grad(set_to_none=True)替代.zero_grad:让内存分配器处理梯度,而不是主动将它们设置为0。
这可以小幅实现加速,但其也有副作用。具体参考:TORCH.OPTIM DOCS

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

将numpy array转化tensor时:torch.tensor()会复制数据,而torch.as_tensor()或者torch.from_numpy不会!

13 关掉不必要的debugging API

PyTorch 提供了许多有用的调试工具,例如 autograd.profilerautograd.grad_checkautograd.anomaly_detection
在需要时使用它们!但在不需要它们时要关闭它们,否则会减慢训练速度。

14 使用gradient clipping

gradient clipping可以大致理解为:gradient = min(gradient, threshold)。其最初是被用于避免 RNN 中的梯度爆炸的问题,但是有一些经验和理论表明其可以加速收敛:torch.nn.utils.clip_grad_norm_()
参考:Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity

gradient clipping对不同模型的加速效果不确定,但它似乎对 RNN、基于 Transformer 和 ResNets 架构以及一系列不同的优化器非常有用。

15 BatchNorm前设置:bias=False

在二维卷积中:torch.nn.Conv2d(..., bias=False, ...)
其作用原理参考:Can not use both bias and batch normalization in convolution layers
该方法可以节约参数,但加速效果比较小。

16 在验证阶段关闭梯度计算

验证阶段设置torch.no_grad()

17 输入归一化&批量标准化

二者都可以参考官方文档。
其解释可以参考:In Machine learning, how does normalization help in convergence of gradient descent?

18 补充:使用 JIT 融合point-wise操作

如果有相邻的逐点操作,可以使用 PyTorch JIT 将它们组合成一个 FusionGroup,然后可以在单个内核上启动,而不是像默认情况下那样在多个内核上启动。这样可以节省一些内存读取和写入。
其使用可以参考:PYTORCH PERFORMANCE TUNING GUIDE、Optimizing CUDA Recurrent Neural Networks with TorchScript

最后

以上就是火星上冬天为你收集整理的加速Pytorch训练的全部内容,希望文章能够帮你解决加速Pytorch训练所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错,欢迎将靠谱客网站推荐给程序员好友。

本图文内容来源于网友提供,作为学习参考使用,或来自网络收集整理,版权属于原作者所有。
点赞(40)

评论列表共有 0 条评论

立即
投稿
返回
顶部