极智AI | 谈谈几种量化感知训练算法

797 阅读7分钟

欢迎关注我的公众号 [极智视界],获取我的更多经验分享

大家好,我是极智视界,本文来介绍一下 谈谈几种量化感知训练算法。

邀您加入我的知识星球「极智视界」,星球内有超多好玩的项目实战源码下载,链接:t.zsxq.com/0aiNxERDq

一般所说的量化算法主要包括后量化 (PTQ) 和 量化感知训练 (QAT),对于工程部署来说,PTQ 往往会是首选,这是因为它的高效 (训练部署解耦)、简单,但是当部署要求提高的时候,比如要求量化误差更加小、比如要求要更低比特的量化时,那就要倾向推 QAT 了。

这里来说说什么是量化感知训练 以及 几种常用的量化感知训练算法。另外关于一些 PTQ 后量化算法的介绍,我之前有写过一些系列文章,感兴趣的同学可以去翻看学习一下。

从以上的讨论其实稍微可以看出 QAT 的几个优势:可以降低量化精度损失、可以达到更加低比特的量化,比如 int4 量化,从而可以进一步轻量化模型、提高推理效率。

量化感知训练的实现方式就是在训练过程中进行模拟量化 或者说 插入伪量化算子 (FakeQuant),这里需要解释一下 FakeQuant 伪量化的概念,它不同于后量化中 Quant 的概念,而更像是后量化中 Quant + DeQuant 的结合。量化感知训练中的 FakeQuant 是将模型中的 weight 和 activation 量化成低比特精度如 int8 后再反量化回 fp32,这个过程会引入量化误差,将其作为训练收敛的指标。所以在整个 finetune 的过程中,模型还是会以 fp32 的精度进行计算,这并不像推理过程中直接纯粹用低比特的精度进行计算,之后再通过反向传播和梯度下降的方式微调模型的权重。

对于量化感知训练过程中量化参数的确定,weight 通常采用绝对最大的值除以 127 的方式确定,而 activate 量化参数则一般根据 QAT 量化算法的不同可能也会有所不同,早期的 QAT 算法采用滑动平均的方式在训练过程中对量化的取值范围进行统计更新,而近些年来主流的 QAT 算法是直接将量化参数定为可学习的参数,在反向传播中通过梯度下降进行更新。这样进行学习的方式会遇到一个比较棘手的问题,就是伪量化过程中 round 截断函数的梯度要怎么计算的问题。这是一个离散函数,是不可导的,换种说法是它的梯度处处为 0,这样在反向传播过程中,权重就无法得到更新。为了解决这个问题,需要想办法对它进行近似操作,一种传统的方式就是采用 Straight-Through Estimator (STE),就是在反向传播过程中,让伪量化算子输出的梯度等于输入的梯度 或者说 使它的梯度为 1。

由上面这种方法引申出了一系列的 QAT 算法,这些算法都可以说是基于 STE 的,比较有代表性的是 DoReFa-Net 量化算法。这个算法将 activation 和 weight 都压缩在 01 之间进行量化,在对激活值进行伪量化前先将激活值截断在 01 之间,这样的做法是基于在一些经典的网络中,比如 resnet 中,大部分激活值都分布在这样的一个范围内。而对于 weight,则是先通过下面这样一个非线性函数对 weight 进行映射,映射到 01 的值域范围,然后对权重值进行量化反量化,最后将结果乘 2 减 1,让值域回到了 -11 的范围 (这也是最开始 tanh(w) 的值域范围)。这里的 quant 和 dequant 在原论文中被表示为 quantizek 的操作,如下,其中 k 指的就是量化的位数。DoReFa-Net 的方式对于其他一些激活值取值分布较广的网络就显得较为死板,所以就延伸了另外一个也是基于 STE 的改进版的量化算法 PACT。

PACT 全称为 Parameterized Clipping Activation,PACT 针对 DoReFa 中 activation 伪量化这一步进行了优化,将截断门限 α 设置为可学习的参数,然后通过梯度下降来寻找更优的门限。因为 α 被视为可学习的参数,所以这部分的导数可以进一步配置为小于 α 的部分为 0,大于 α 的部分为 1。可以看到,PACT 这个函数其实和 ReLU 很像,都是对大于 0 的部分做一个截断操作,其实 PACT 的原本思想也是想去替代掉 ReLU。

但问题是并非所有的网络都是用 ReLU 这个激活函数的,所以为了让 PACT 的应用范围更广,在 PaddlePaddle 中也对其进行了改进,对于大于 0 和小于 0 的激活值都做了相同的限制,使它能在某些情况下得到更好的量化范围,以降低量化损失。

除此之外还有个常被用到 QAT 量化算法是 LSQ,全称为 Learned Step Size Quantization。与 PACT 相似的地方是,它也是通过训练来确定量化参数的,但不同的是,PACT 学习的是截断门限,而 LSQ 是直接将 Scale 定为一个可直接学习的参数,计算公式如下。这样的公式其实就是在对称量化的基础上加上了一个 Scale,完成了一个量化反量化的过程。因为需要反向传播计算梯度,所以需要对下面的公式进行求导。

直接求导的公式如下,

中间部分用 STE 进行简化,如下,

这里也可以看出,LSQ 虽然也采用了 STE 的直通估计,但它在截取范围内还是存在梯度计算的,而且为了使得 Scale 的学习更加稳定,LSQ 还为其梯度加上了一个缩放系数,如下,这个系数主要是由 Tensor 中的元素数量决定。

另外既然是可学习的参数,就需要为其设置一个初始值,PACT 中 α 初始值一般会直接设置为 6,而 LSQ 原文中 Scale 的初始值是计算出来的,计算方式如下。不过在实践过程中,其实更多的是通过统计激活值分布,用 EMMAX、KLD 或 MSE 等 PTQ 的方式来计算初始的 Scale 值。

虽然使用 STE 可以简单粗暴地解决量化函数没有梯度的问题,但是也有研究者认为 STE 估算出来的梯度并不是特别准确,他们对 round 函数采取了近似方式,例如商汤在 19 年提出了 DSQ (Differentiable Soft Quantization) 量化算法,利用一个可导的双曲正切函数去近似阶梯函数,从而代替 round 的操作,在 QAT 中进行伪量化。然后又在这样的函数上加了一个可学习的参数来控制曲线的形状,随着训练的推进,这个参数的更新会使得函数的形状越来越接近真实的阶梯函数。

另外,在深度学习模型中,在推理阶段一般会做算子融合,比较典型的是 Conv + BN + ReLU 的融合,但这样就造成了一个问题,就是推理阶段的路径和训练阶段的路径不一样。针对这种算子融合的情况,在 QAT 中,一般会先把 BN 中的 γ、β、mean、std 放进 Conv weight 和 bias 中,再输入伪量化算子。这样的做法下,就不会在训练的过程中,更新 BN 的参数。当然,其他算子融合的情况,QAT 也是需要做相应的优化调整的,这部分其实是属于比较复杂的。

好了,以上分享了 谈谈几种量化感知训练算法,希望我的分享能对你的学习有一点帮助。



 【公众号传送】

《极智AI | 谈谈几种量化感知训练算法》


畅享人工智能的科技魅力,让好玩的AI项目不难玩。邀请您加入我的知识星球, 星球内我精心整备了大量好玩的AI项目,皆以工程源码形式开放使用,涵盖人脸、检测、分割、多模态、AIGC、自动驾驶、工业等。不敢说会对你学习有所帮助,但一定非常好玩,并持续更新更加有趣的项目。 t.zsxq.com/0aiNxERDq

logo_show.gif