用上Pytorch Lightning的这六招,深度学习pipeline提速10倍!

金磊 发自 凹非寺
量子位 报道 | 公众号 QbitAI

面对数以亿计的图片数据,到底该用什么样的方法才能快速搞实验?

这样的问题,或许在做机器学习研究的你,也会经常遇到。

而就在最近,一个国外小哥就提出了一种建议:

在Pytorch lightning基础上,让深度学习pipeline速度提升10倍

127b59d5d3deda5b8b9e65a78467096d.gif

用他自己的话来说就是——“爬楼时像给了你一个电梯”。

这般“酸爽”,到底是如何做到的呢?

7668cc4bf1a4734f441d6493be3e92ab.png

优化机器学习pipeline,很重要

无论你是身处学术界还是工业界,时间资源等各种因素,往往会成为你在搞实验的枷锁

尤其是随着数据集规模和机器学习模型,变得越发庞大和复杂,让实验变得既费时又耗力。

c6ef1965910b55525e9500d799fd9962.png

提速这件事,就变得至关重要。

例如在2012年的时候,训练一个AlexNet,要花上5到6天的时间。

而现如今,只需要短短几分钟就可以在更大的数据集上训练更大的图像模型。

这位小哥认为,从某种角度上来说,这是得益于各种各样的“利器”的出现。

例如Pytorch Lingtning,就是其中一种。

于是,他便“死磕”pipeline,总结了六种“闪电加速”实验周期的方法。

并行数据加载

数据加载和增强(augmentation)往往被认为是训练pipeline时的瓶颈之一。

一个典型的数据pipeline包含以下步骤:

  • 从磁盘加载数据

  • 在运行过程中创建随机增强

  • 将每个样本分批整理

在这个过程中,倒是可以用多个CPU进程并行加载数据来优化。

但与此同时,还可以通过下面的操作来加速这一过程:

1、将DataLoader中的num_workers参数设置为CPU的数量。

2、当与GPU一起工作时,将DataLoader中的pin_memory参数设置为True。这可以将数据分配到页锁定的内存中,从而加快数据传输到GPU的速度。

使用分布式数据并行的多GPU训练

1d4bf9717823625ce9a8611de23b3ba3.png

与CPU相比,GPU已经大大加速了训练和推理时间。

但有没有比一个GPU更好的方法?或许答案就是:

多个GPU!

在PyTorch中,有几种范式可以用多个GPU训练你的模型。

两个比较常见的范式是 “DataParallel ”和 “DistributedDataParallel”。

而小哥采用的方法是后者,因为他认为这是一种更可扩展的方法。

但在PyTorch(以及其他平台)中修改训练pipeline并非易事。

必须考虑以分布式方式加载数据以及权重、梯度和指标的同步等问题。

不过,有了PyTorch Lightning,就可以非常容易地在多个GPU上训练PyTorch模型,还是几乎不需要修改代码的那种!

a9c0754cd171560978396161065cf792.png

混合精度

在默认情况下,输入张量以及模型权重是以单精度(float32)定义的。

然而,某些数学运算可以用半精度(float16)进行。

这样一来,就可以显著提升速度,并降低了模型的内存带宽,还不会牺牲模型的性能。

通过在PyTorch Lightning中设置混合精度标志(flag),它会在可能的情况下自动使用半精度,而在其他地方保留单精度。

通过最小的代码修改,模型训练的速度可以提升1.5至2倍。

42fc5bf9ba1593fbdc44e13fcd1d67ec.png

早停法

当我们训练深度学习神经网络的时候,通常希望能获得最好的泛化性能。

但是所有的标准深度学习神经网络结构,比如全连接多层感知机都很容易过拟合。

当网络在训练集上表现越来越好,错误率越来越低的时候,实际上在某一刻,它在测试集的表现已经开始变差。

因此,早停法 (Early Stopping)便在训练过程中加入了进来。

具体来说,就是当验证损失在预设的评估次数(在小哥的例子中是10次评估)后停止训练。

这样一来,不仅防止了过拟合的现象,而且还可以在几十个 epoch内找到最佳模型。

02434550483b4ca880cd70f9dcfb57da.png

Sharded Training

Sharded Training是基于微软的ZeRO研究和DeepSpeed库。

它显著的效果,就是让训练大模型变得可扩展和容易。

否则,这些模型就不适合在单个GPU上使用了。

而在Pytorch Lightning的1.2版本中,便加入了对Shared Training的支持。

虽然在小哥的实验过程中,并没有看到训练时间或内存占用方面有任何改善。

但他认为,这种方法在其它实验中可能会提供帮助,尤其是在不使用单一GPU的大模型方面。

c4cd72b1d49e3860dd2e88bb7d8c0525.png

模型评估和推理中的优化

在模型评估和推理期间,梯度不需要用于模型的前向传递。

因此,可以将评估代码包裹在一个torch.no_grad上下文管理器中。

这可以防止在前向传递过程中的存储梯度,从而减少内存占用。

如此一来,就可以将更大的batch送入模型,让评估和推理变得更快。

效果如何?

介绍了这么多,你肯定想知道上述这些方法,具体起到了怎样的作用。

小哥为此做了一张表格,详解了方法的加速效果。

4c84df9548e3711e7837676f76746ccc.png

那么这些方法,是否对在做机器学习实验的你有所帮助呢?

快去试试吧~

参考链接:

https://devblog.pytorchlightning.ai/how-we-used-pytorch-lightning-to-make-our-deep-learning-pipeline-10x-faster-731bd7ad318a

  • 0
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值