用上PytorchLightning的这六招，深度学习pipeline提速10倍！

点击上方“视学算法”&＃xff0c;选择加"星标"或“置顶”

重磅干货&＃xff0c;第一时间送达

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

面对数以亿计的图片数据&＃xff0c;到底该用什么样的方法才能快速搞实验&＃xff1f;

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

而就在最近&＃xff0c;一个国外小哥就提出了一种建议&＃xff1a;

在Pytorch lightning基础上&＃xff0c;让深度学习pipeline速度提升10倍&＃xff01;

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

这般“酸爽”&＃xff0c;到底是如何做到的呢&＃xff1f;

优化机器学习pipeline&＃xff0c;很重要

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

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

提速这件事&＃xff0c;就变得至关重要。

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

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

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

例如Pytorch Lingtning&＃xff0c;就是其中一种。

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

并行数据加载

数据加载和增强&＃xff08;augmentation&＃xff09;往往被认为是训练pipeline时的瓶颈之一。

一个典型的数据pipeline包含以下步骤&＃xff1a;

• 从磁盘加载数据

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

• 将每个样本分批整理
  

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

但与此同时&＃xff0c;还可以通过下面的操作来加速这一过程&＃xff1a;

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

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

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

与CPU相比&＃xff0c;GPU已经大大加速了训练和推理时间。

但有没有比一个GPU更好的方法&＃xff1f;或许答案就是&＃xff1a;

多个GPU&＃xff01;

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

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

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

但在PyTorch&＃xff08;以及其他平台&＃xff09;中修改训练pipeline并非易事。

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

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

混合精度

在默认情况下&＃xff0c;输入张量以及模型权重是以单精度&＃xff08;float32&＃xff09;定义的。

然而&＃xff0c;某些数学运算可以用半精度&＃xff08;float16&＃xff09;进行。

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

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

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

早停法

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

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

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

因此&＃xff0c;早停法 &＃xff08;Early Stopping&＃xff09;便在训练过程中加入了进来。

具体来说&＃xff0c;就是当验证损失在预设的评估次数&＃xff08;在小哥的例子中是10次评估&＃xff09;后停止训练。

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

Sharded Training

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

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

否则&＃xff0c;这些模型就不适合在单个GPU上使用了。

而在Pytorch Lightning的1.2版本中&＃xff0c;便加入了对Shared Training的支持。

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

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

模型评估和推理中的优化

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

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

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

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

效果如何&＃xff1f;

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

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

那么这些方法&＃xff0c;是否对在做机器学习实验的你有所帮助呢&＃xff1f;

快去试试吧~

参考链接&＃xff1a;

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

— 完 —

本文系网易新闻•网易号特色内容激励计划签约账号【量子位】原创内容&＃xff0c;未经账号授权&＃xff0c;禁止随意转载。

点个在看 paper不断&＃xff01;