简介：本文以近期风行的 YOLOX[8]指标检测模型为例，介绍量化感知训练的原理流程，探讨如何实现精度无损的实践经验，并展现了量化后的模型可能做到精度不低于原始浮点模型，模型压缩 4X、推理减速最高 2.3X 的优化成果。
1. 概述
对深度学习模型进行低比特量化，能够无效地升高模型部署时在存储、计算、通信上的开销，是一种常见的模型压缩和推理优化技术。然而，模型量化在理论利用中依然存在不少挑战，最为常见的问题就是模型精度降落（如无非凡阐明，本文中“模型精度”是指准确率等模型利用于具体任务的成果指标）。以计算机视觉畛域为例，在指标检测、图像宰割等简单工作上，量化带来的精度降落更为显著。
通过在模型训练阶段引入量化相干束缚，即量化感知训练（Quantization-aware training，QAT），可能更好地解决模型量化的精度问题。本文以近期风行的 YOLOX[8]指标检测模型为例，介绍量化感知训练的原理流程，探讨如何实现精度无损的实践经验，并展现了量化后的模型可能做到精度不低于原始浮点模型，模型压缩 4X、推理减速最高 2.3X 的优化成果。
2. 量化原理
在数字信号处理畛域，量化是指将信号的间断取值（或者大量可能的离散取值）近似为无限多个（或较少的）离散值的过程。具体到深度学习畛域，模型量化是指将浮点激活值或权重（通常以 32 比特浮点数示意）近似为低比特的整数（16 比特或 8 比特），进而在低比特的示意下实现计算的过程。通常而言，模型量化能够压缩模型参数，进而升高模型存储开销；并且通过升高访存和无效利用低比特计算指令等，可能获得推理速度的晋升，这对于在资源受限设施上部署模型尤为重要。
给定浮点类型的值，能够通过如下公式将它转化成 8 比特量化值：
其中，示意量化的 scale，与别离示意量化值域的最小值与最大值，示意输出浮点值，示意量化后的值。量化值转化为浮点值只需执行反操作即可：
进一步的，在将输出数据和权重进行量化后，咱们就可能将神经网络的常见操作转换为量化操作。以卷积操作为例，其量化版本典型的计算流程如图 1 所示：
图 1 典型的量化卷积算子计算流程图
- 权重与输出先量化成 8bit，进行卷积操作，用 32bit 来存储两头后果；
- bias 量化为 32bit，与进行相加为；
- 利用、以及的量化 scale 将 32bit 的转化为 8bit 的；
- 如果该层的下一层也是量化 OP，则可间接输入给下一层；如果是非量化 OP，则将反量化为浮点值后，再输入给下一层。
从上述量化计算的原理可能容易看出，将浮点数转化为低比特整数进行计算会不可避免地引入误差，神经网络模型中每层量化计算的误差会累积为模型整体精度的误差。常见的训练后量化（Post Training Quantization，PTQ）计划中，通过统计在典型输出数据状况下，待量化变量的数值散布，来抉择适合的量化参数（scale，zero point 等），将因量化而引入的信息损失升高到最小。
然而 PTQ 计划往往还是无奈实现精度无损的模型量化，为了进一步升高量化带来的精度降落，咱们能够采纳量化感知训练的计划，在训练的计算图中引入伪量化的操作，通过微调训练 (finetuning) 让模型权重“适应”量化引入的误差，以实现更好的、甚至无损的量化模型精度。
3. YOLOX 量化训练
咱们以 YOLOX- s 指标检测模型 (GitHub repo[1]) 为例，应用公开的预训练模型参数，在 COCO2017 数据集上进行量化训练试验。量化训练算法抉择 LSQ[2,3]，该系列算法利用梯度来更新量化的 scale 与 zero_point，在不须要精密调节参数的状况下可能取得较好的性能。为了通过量化训练取得更好的量化模型精度，咱们须要重点关注如下几点设置：
3.1 与部署后端相匹配的量化形式
不同的部署后端，可能采纳不必的量化计算实现形式，须要匹配训练和部署阶段的量化形式以防止引入额定误差。以 PyTorch[7]默认的 CPU 后端为例，根本的量化形式为
- weight：per-channel，int8，对称量化
- activation：per-tensor，uint8，非对称量化
- 量化对象：所有的 Conv2d
以挪动端框架 MNN 为例，根本的量化形式为：
- weight：per-channel，int8，对称量化
- activation：per-tensor，int8，对称量化
- 量化对象：所有的 Conv2d
通常在具体推理框架上部署量化模型时，还会对相似 conv-bn、conv-relu/relu6、conv-bn-relu/relu6 这样的卷积层进行算子交融，因而须要设置为量化交融后的 activation。YOLOX 中的 SiLU 激活函数通常不会被交融，output activation 的量化仅须要设置到 bn 的输入，而不是 SiLU 的输入。训练时一个典型的卷积层量化地位如图 2 所示。
图 2 量化地位示意图
同时，QAT 时所有的 BN 全副 fold 入相应的卷积中，试验采取了文献 [6] 中的 fold 策略。因为模仿量化的引入可能会使得 BN 层的 running_mean 与 running_var 不稳固，从而导致量化训练无奈收敛。因而，咱们在量化训练中固定 BN 的 running_mean 与 running_var。
此外，特定部署后端的实现可能须要将 activation 量化为 7bit 数值范畴，以避免计算溢出。带有 avx512_vnni 指令集的 CPU 型号上，则没有相应要求。
3.2 量化训练参数初始化
为了防止导致训练 loss 不稳固甚至训练发散的状况，咱们利用训练后量化（Post training quantization，PTQ）所得的量化参数来初始化 LSQ 量化训练中 activation scale 参数。根本的步骤是：
- 选取几十张典型图片，在预训练模型上进行推理，并收集统计各个待量化的 activation 信息。
- 应用如 MSE、KL 散度等 metric 来计算各 activation 的最佳 scale。
- 根本的实现形式能够应用 PyTorch 的 HistogramObserver 计算 activation scale & zero point，PerChannelMinMaxObserver 计算 weight scale。
3.3 训练超参数
在 QAT 阶段，咱们应用曾经收敛的预训练模型权重，并且应用 PTQ 量化参数进行 QAT 初始化。这种状况下，模型参数已靠近收敛，因而咱们将整体训练的超参数和原始 YOLOX 训练保持一致，将学习率设置为原始训练收敛阶段的学习率，即 5e-4。
3.4 特定后端算子实现的计算误差
因为数值示意精度问题，round 操作的后果在训练框架与推理后端上可能并不是雷同的，例如：
import torch
torch.tensor(2.5).cuda().round() # 输入 tensor(2., device=’cuda:0′)
torch.tensor(3.5).cuda().round() # 输入 tensor(4., device=’cuda:0′)
2.5 和 3.5 的四舍五入行为在 PyTorch 上并不是雷同的，而该行为差别在常见的部署后端框架上可能不存在。这种问题会导致 QAT 模仿精度和后端理论运行时的精度存在显著差别，因而须要在量化训练阶段进行修改。例如针对 MNN 后端，咱们能够采取如下操作来防止这个差别。
def round_pass(x):
"""A simple way to achieve STE operation."""
# y = torch.round(x) # for PyTorch backend
y = (x + torch.sign(x) * 1e-6).round() # for mnn backend, round 前增加一个小的数值
y_grad = x
return (y - y_grad).detach() + y_grad
4. 试验后果与剖析
4.1 精度和减速成果
依照上述形式在 YOLOX- s 模型上进行量化训练，应用 COCO2017 validation set 进行精度验证，后果如表 1 所示。两种后端上实在量化的模型的精度性能均和浮点模型齐平。
表 1 浮点模型、QAT 模型及后端实在量化模型的精度比照
速度试验中，咱们选取 PyTorch[7]后端及 x86 测试平台进行测试，测试的图片分辨率为 1x3x640x640。不同资源数下的后果如表 2 所示。在精度无损的前提下，量化模型的推理速度最高能够晋升 2.35x，同时模型尺寸为原来的 1 /4。
表 2 浮点模型、QAT 模型及后端实在量化模型的速度比照
4.2 量化参数初始化的影响
LSQ 及 LSQ+ 论文中提出了相应的量化信息初始化办法，理论应用中会发现该种初始化办法对于学习率的设定会比拟敏感。一方面，若量化参数的初始值间隔收敛值较远，须要通过设置较大的学习率来训练他们。另一方面，较大的学习率会使得模型自身参数跳出已收敛的状态，使得模型进入“重训练”的状态，不确定性减少。而用 PTQ 来初始化量化参数，能够使得参数处于较优的初始状态，利用浮点模型收敛时的学习率去进行 finetune 即可取得较好的性能。固定学习率为 4e- 5 的前提下，PTQ 及 LSQ 初始化的训练后果与曲线如图 4 所示。LSQ 初始化的训练曲线在开启 mossac 数据加强时是逐步向下的，而 PTQ 初始化是逐步向上的。
图 4 固定 finetune 学习率下，不同初始化办法的训练曲线
4.3 训练超参数的影响
学习率的设定会间接影响到最终 QAT 模型的性能。以原始模型训练收敛阶段的学习率（5e-4）为基准，如果 QAT 阶段应用雷同的学习率，QAT 初期的模型精度会逐步降落，如图 4 中 PTQ 初始化红色曲线训练晚期所示，然而最终精度会晋升至原始非量化模型的程度。如果 QAT 应用更低学习率（例如 5e-4），模型精度会相比于 PTQ 初始化状态逐步回升，然而最终精度晋升不大：
上述景象的一种可能起因是，在小学习率下模型权重、量化 scale 根本不变。实际上是基于 PTQ 的初始解空间持续向部分更好的收敛点凑近，跳出部分解空间的可能性较低。因而，在 PTQ 初始化的前提下，学习率能够间接设置成浮点模型训练收敛阶段的值。
4.4 训练轮数的抉择
上述 QAT 的后果是训练 300 epochs 后的模型，缩小训练 epoch 数量的后果，如下所示：
能够看出随着训练轮数的增大，QAT 的后果会更好。QAT 训练轮数较低，后果会不如间接用小学习率进行 finetune。咱们能够得出经验性的 trade-off：如果计算资源短缺，能够抉择训练更长的工夫来取得更好的性能。如果计算资源较少，则能够用小学习率训练较短的工夫，以取得高于 PTQ 的性能。
4.5 修改特定算子计算误差的影响
特定算子在训练框架与后端框架中的行为会有轻微的差异，这会导致量化训练精度与理论量化模型的精度产生较大的差别。以 MNN[5]为例，在 YOLOX 模型中存在如下两类 OP 会导致该景象。
修改 round 操作的影响
是否对 round 进行修改的后果如表 3 所示。从训练角度而言，round 通过修改后的性能会略好于未修改的。同时如果在训练时 round 的行为与后端不统一的话，可能会导致实在量化模型的精度产生较大的变动。
表 3 round 修改对于量化训练的影响
Sigmoid 疾速实现引入的误差
为了晋升指数计算的速度，后端框架通常会采取一些疾速近似计算。这对于浮点模型而言，通常不会引入较大的误差。然而对于量化模型而言，这个误差可能因为量化操作而被放大。
如图 3 所示，对于 YOLOX 的次要 pattern（Conv -> SiLU -> Conv），前一层 Conv 的输入通过 SiLU 函数后，疾速近似计算引入的误差会被后一层卷积输出处的量化操作（除以 scale）而缩放。scale 越小，缩放的水平越大。以 YOLOX- s 为例，是否对指数计算进行近似的量化模型精度如表 4 所示。
图 3 YOLOX 量化 pattern
表 4 指数计算的近似带来的模型性能误差
5. 总结
本文通过在 YOLOX 指标检测模型上的量化实际，验证了通过量化感知训练 (QAT) 可能在精度无损的状况下，取得显著的模型压缩和推理减速。咱们对量化精度误差因素进行了具体分析，指出了解决精度问题的一系列实际伎俩，并通过试验验证了成果，这能够作为咱们在理论利用模型量化时的教训参考。尽管量化相干办法曾经被大量钻研，然而在理论简单工作中利用量化依然面临不少挑战，真正将量化压缩落地，更须要通过模型和零碎两方面的协同。此外，还有更多量化训练相干的技术计划（如混合精度量化、更低比特量化等）值得摸索和欠缺。
对于咱们
本文中对于量化训练实际工作由阿里云 -PAI 模型压缩团队和微软 NNI 团队单干实现，也感激 MNN 团队的技术支持。更多模型压缩的算法实现可参考 NNI（GitHub repo[4]），更多模型推理优化技术计划可见阿里云 PAI-Blade。
参考文献 & 代码仓库
[1] https://github.com/Megvii-Bas…
[2] Esser S K, McKinstry J L, Bablani D, et al. Learned step size quantization[J]. arXiv preprint arXiv:1902.08153, 2019.
[3] Bhalgat Y, Lee J, Nagel M, et al. Lsq+: Improving low-bit quantization through learnable offsets and better initialization[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops. 2020: 696-697.
[4] https://github.com/microsoft/nni
[5] Jiang X, Wang H, Chen Y, et al. Mnn: A universal and efficient inference engine[J]. arXiv preprint arXiv:2002.12418, 2020.
[6] Jacob B, Kligys S, Chen B, et al. Quantization and training of neural networks for efficient integer-arithmetic-only inference[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2018: 2704-2713.
[7] Paszke A, Gross S, Massa F, et al. Pytorch: An imperative style, high-performance deep learning library[J]. Advances in neural information processing systems, 2019, 32: 8026-8037.
[8] Ge Z, Liu S, Wang F, et al. Yolox: Exceeding yolo series in 2021[J]. arXiv preprint arXiv:2107.08430, 2021.
原文链接
本文为阿里云原创内容，未经容许不得转载。