1. 概述
对深度学习模型进行低比特量化,可以有效地降低模型部署时在存储、计算、通信上的开销,是一种常见的模型压缩和推理优化技术。然而,模型量化在实际应用中仍然存在不少挑战,最为常见的问题就是模型精度下降(如无特殊说明,本文中“模型精度”是指准确率等模型应用于具体任务的效果指标)。以计算机视觉领域为例,在目标检测、图像分割等复杂任务上,量化带来的精度下降更为明显。
通过在模型训练阶段引入量化相关约束,即量化感知训练(Quantization-aware training,QAT),能够更好地解决模型量化的精度问题。本文以近期流行的YOLOX[8]目标检测模型为例,介绍量化感知训练的原理流程,讨论如何实现精度无损的实践经验,并展示了量化后的模型能够做到精度不低于原始浮点模型,模型压缩4X、推理加速最高2.3X的优化效果。
2. 量化原理
在数字信号处理领域,量化是指将信号的连续取值(或者大量可能的离散取值)近似为有限多个(或较少的)离散值的过程。具体到深度学习领域,模型量化是指将浮点激活值或权重(通常以32比特浮点数表示)近似为低比特的整数(16比特或8比特),进而在低比特的表示下完成计算的过程。通常而言,模型量化可以压缩模型参数,进而降低模型存储开销;并且通过降低访存和有效利用低比特计算指令等,能够取得推理速度的提升,这对于在资源受限设备上部署模型尤为重要。
进一步的,在将输入数据和权重进行量化后,我们就能够将神经网络的常见操作转换为量化操作。以卷积操作为例,其量化版本典型的计算流程如图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后端为例,基本的量化方式为