基于PyTorch智能混合精度库对ResNet50模型启动优化训练

本文将经过一个实战案例来展现如何借助于PyTorch智能混合精度库对ResNet50模型启动优化，而后借助少许几行代码即可取得超越两倍速度的模型训练效率。

简介

你能否曾宿愿你的深度学习模型运转得更快？

但是，咱们该选哪一种模型呢？

、 HuggingFace 和曾经为模型训练的性能调优提供了很好的参考，包括异步数据加载、缓冲区审核点、散布式数据并行化和智能混合精度等等。

在这篇文章中，我将专一引见智能混合精度技术。首先，我将简明引见Nvidia的张量核设计；而后，咱们一同讨论宣布在ICLR 2018上的开创性上班——“混合精度训练”关系论文；最后，我将引见一个在FashionMNIST数据集上训练ResNet50模型的繁难示例。经过这个示例，咱们来展现如何在加载双倍批量数据的同时将训练速度提高两倍，而这一结果却只要要额外编写三行代码。

配件基础——Nvidia张量核

首先，让咱们回忆一下GPU设计的一些基本原理。英伟达GPU最受欢迎的商业产品之一是Volta系列，例如基于GV100 GPU设计的V100 GPU。因此，咱们将围绕上方的GV100架构启动讨论。

关于GV100架构来说，流式多处置器（SM）是计算的外围设计。每个GPU蕴含6个GPU处置集群（GPC）和S84 SM（或V100的80 SM）。其全体设计如下图所示：

关于每个SM，它蕴含两种类型的外围：CUDA外围和张量核。CUDA外围是Nvidia于2006年推出的原始设计，是CUDA平台的关键组成局部。CUDA外围可分为三种类型：FP64外围/单元、FP32外围/单元和Int32外围/单元。每个GV100 SM蕴含32个FP64外围、64个FP32外围和64个Int32外围。Volta/Turing（2017）系列GPU中引入了张量核，以便与之前的Pascal（2016）系列分别。GV100上的每个SM蕴含8个张量核。 链接 处给出了V100 GPU的完整详细信息列表。上方详细引见SM设计。

为什么选用张量核？Nvidia张量核专门用于口头通用矩阵乘法（GEMM）和半精度矩阵乘法和累加（HMMA）操作。简而言之，GEMM以A*B+C的格局口头矩阵运算，HMMA将运算转换为半精度格局。无关这方面的更详细的讨论可以在 链接 处找到。由于深度学习触及MMA；所以，张量（Tensor）外围在当今的模型训练和减速计算中至关关键。

当然，当切换到混合精度训练时，请务必审核你经常使用的GPU的规格。只要最新的GPU系列允许张量核，混合精度训练只能在这些机器上经常使用。

数据格局基础——单精度（FP32）与半精度（FP16）

如今，让咱们细心看看FP32和FP16格局。FP32和FP16是IEEE格局，经常使用32位二进制存储和16位二进制存储示意浮点数。这两种格局都包括三个局部：a）符号位；b）指数位；c）尾数位。FP32和FP16调配给指数和尾数的比特数不同，这造成了不同的值范围和精度。

如何将FP16和FP32转换为实在的值呢？依据IEEE-754规范，FP32的十进制值=（-1）^（符号）×2^（十进制指数-127）×（隐式前导1+十进制尾数），其中127是有偏向的指数值。关于FP16，公式变为（-1）^（符号）×2^（十进制指数-15）×（隐式前导1+十进制尾数），其中15是相应的有偏指数值。你可以在 链接 处检查有偏指数值的更多详细信息。

从这个意义上讲，FP32的取值范围约为[-2¹²S，2¹²83;]~[-1.7*1e38，1.7*1e38]，FP16的取值范围大概为[-2⁵，2'8309]=[-32768，32768]。请留意，FP32的十进制指数在0到255之间，咱们扫除了最大值0xFF，由于它示意NAN。这就解释了为什么最大的十进制指数是254–127=127。当然，相似的规定也实用于FP16。

关于精度方面，请留意指数和尾数都有助于精度限度（也称为非规范化，请参阅 链接 处的详细讨论）。因此，FP32可以示意高达2^（-23）*2^（-126）=2^（-149）的精度，FP16可以示意高至2^（10）*2^。

FP32和FP16示意之间的差异带来了混合精度训练的关键疑问，由于深度学习模型的不同层/操作对值范围和精度或许不敏感或许敏感，所以须要独自处置。

混合精度训练

前面，咱们曾经学习了MMA的配件基础常识、张量核的概念以及FP32和FP16之间的关键区别。接上去，咱们可以进一步讨论混合精度训练的细节。

混合精度训练的想法最早是在2018年ICLR论文 《混合精度训练》（Mixed Precision Training） 中提出的。该论文在训练环节中将深度学习模型转换为半精度浮点，而不会损失模型精度或修正超参数。如上所述，由于FP32和FP16之间的关键区别在于值范围和精度，该论文详细讨论了FP16为什么会造成梯度隐没，以及如何经过损失缩放来处置这个疑问。此外，该论文还提出了经常使用FP32主权重拷贝和经常使用FP32启动归约和向量点积攒加等特定操作的技巧。

损失缩放（Loss scaling）。本文给出了一个经常使用FP32精度训练Multibox SSD探测器网络的示例，如下所示。假设不启动任何缩放，FP16梯度的指数范围≥2^（-24），以下一切值都将变为零，这与FP32相比是不够的。但是，经过试验，将梯度繁难地缩放2³=8倍，可以使半精度训练精度复原到与FP32相婚配的水平。从这个意义上讲，作者以为[2^（-27），2^（-24）]之间的额外百分之几的梯度在训练环节中依然很关键，而低于2^（-27）的值并不关键。

处置这种规模差异的方法是借助损失缩放的方法。依据链式规律，缩放损失将确保相反的量将缩放一切梯度。但是请留意，在最终权重更新之前，须要敞开缩放梯度。

智能混合精度训练

Nvidia公司首先开发了名为APEX的PyTorch裁减智能混合精度训练，而后被PyTorch、TensorFlow、MXNet等干流框架宽泛驳回。请参阅 链接 处的Nvidia文档。为了繁难起见，咱们只引见PyTorch框架中的智能混合精度库：

amp库可以智能处置大少数混合精度训练技术，如FP32主权重复制。开发人员只要要启动操作数智能转换和梯度/损失缩放。

操作数智能转换：虽然咱们提到张量核可以大大提高GEMM操作的性能，但某些操作不适宜半精度示意。

amp库给出了一个合乎半精度条件的 CUDA操作列表 。amp.autocast齐全涵盖了大少数矩阵乘法、卷积和线性激活运算；但是，关于归约/求和、softmax和损失计算等，这些计算依然在FP32中口头，由于它们对数据范围和精度更敏感。

梯度/损失缩放：amp库提供了 智能梯度缩放技术 ；因此，用户在训练环节中不用手动调整缩放。 链接 处可以找到缩放因子的更详细的算法。

一旦缩放了梯度，就须要在启动梯度剪裁和正则化之前将其增加。更多细节可以在 链接 处找到。

FashionMNIST训练示例

torch.amp库相对易于经常使用，只要要三行代码即可将训练速度提高2倍。

首先，咱们从一个十分繁难的义务开局，经常使用FP32在FashionMNIST数据集（MIT容许证）上训练ResNet50模型；咱们可以看到10个世代的训练时期为333秒：

ResNet50模型在数据集FashionMNIST上的训练

小于2**（-24）的梯度与总梯度之比。咱们可以看到，FP16将使总梯度的近1/4变为零

评价结果

如今，咱们经常使用amp库。amp库只要要三行额外的代码启动混合精度训练。咱们可以看到训练在141秒内成功，比FP32训练快2.36倍，同时到达了相反的准确度、召回率和F1分数。

scaler = torch.cuda.amp.GradScaler()#开局训练的代码# ...with torch.autocast(device_type="cuda"):#训练代码#封装损失与优化器scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()

经常使用amp库启动训练

训练时期的缩放因子（缩放因子仅在第一步出现变动，坚持不变。）

最终结果与FP32训练结果比拟

上方代码的Github链接：

总结

混合精度训练是减速深度学习模型训练的一种十分有价值的技术。它不只放慢了浮点运算的速度，还节俭了GPU内存，由于训练批次可以转换为FP16，从而节俭了一半的GPU内存。另外，借助于PyTorch框架中的amp库，额外的代码可以增加到仅仅三行，由于权重复制、损失缩放、操作类型转换等计算都是由该库外部处置的。

须要留意的是，假设模型权严重小远大于数据批次的话，混合精度训练并不能真正处置GPU内存疑问。首先，只要模型的某些层被转换成FP16，而其他层仍在FP32中计算；其次，权重更新须要FP32复制，这依然须要占用少量的GPU内存；第三，Adam等优化器的参数在训练环节中占用了少量GPU内存，而混合精度训练使优化器参数坚持不变。从这个意义上说，须要更先进的技术，如DeepSpeed的ZERO算法。

译者引见

朱先忠，社区编辑，专家博客、讲师，潍坊一所高校计算机老师，自在编程界老兵一枚。

原文题目： The Mystery Behind the PyTorch Automatic Mixed Precision Library ，作者：Mengliu Zhao