后训练量化

原文链接：https://www.yuque.com/yahei/hey-yahei/quantization-post_training
欢迎转载&引用，但烦请注明出处~

按照是否需要训练划分，量化通常可以分为从头训练（train from scratch）、重训练（retrain）、后训练（post-training）三种，本文主要介绍几种后训练量化的方案，并以线性均匀分布的定点量化为例。


后训练量化指的是，对预训练后的网络选择合适的量化操作和校准操作，以实现量化损失的最小化，该过程不需要训练，通常不直接更新权重原始数值而是选用合适的量化参数。当然也有一些利用校准数据集适当更新的权重原始数值的技术出现，这类优化方法有点“训练”的意味，但通常所需的“训练”样本非常少，前传（甚至有的需要局部反传）的次数小于训练时一个epoch的所需次数。

常见的优化目标

通常特征图、权重都可以转换成向量的形式来进行优化，所以可以不失一般性地假设我们的优化对象就是向量。

误差函数

误差函数通常就是距离函数，常见的距离函数有三种——

1. L2距离，用于衡量两向量在空间位置上的差异
   
   很多情况下我们并不关注准确的L2距离，只需要比较相对大小，所以很多时候L2距离都指的是平方和误差，也即
   
   后文如不特殊说明，L2距离均指的是平方和误差。

L2距离也可以有其他变种，比如加权版本的L2距离——

从裁剪的经验上看，我们往往认为绝对值较大的权重会比较重要，那么很自然地想到我们可以为L2距离进行加权，增加大权重在量化后误差在整体误差上的重要性。从实践上看，加权版本的L2距离在较高比特的时候似乎没有明显的优势，甚至可能带来更差的结果；但在较低比特量化的应用场景中，似乎能带来更好的结果（相对明显的提升）。

L2距离是一种相对直观的选择，也是最常用的选择。除此之外，L2距离有一些很好的特性，比如它很方便求导，可以很容易推导出它的迭代公式，进行交替优化。

1. KL距离，用于衡量两向量在分布上的差异
   为向量和统计直方图（注意bin数要一致），然后进行归一化操作，使所有bin上的数值相加后为1，也即得到概率分布图和，那么向量和的KL距离为
   
   KL其实不是严格的“距离”，它没有对称性（也即），对应地产生了一些KL的变种，如JS距离、推土机距离，也可以用作误差函数；

衡量分布差异的误差函数对数据分布比较敏感，通常需要较大的数据量来保证收集的分布与真实分布足够接近，而L2距离和余弦距离则没有这样的要求。量化的实际应用最早算是英伟达的TensorRT推动起来的，其校准策略就是最小化KL距离，因此多数校准策略也都模仿它采用了KL距离。

1. 余弦距离，用于衡量两向量在方向上的差异
   
   余弦距离作为误差函数我没用过，确实也有些论文（比如EasyQuant）利用余弦距离作为目标进行优化，但与L2距离和KL距离不同，余弦距离取值范围为 [-1, 1] ，越接近1表示误差越小，所以优化时是最大化余弦距离而非最小化。

优化对象

1. 直接最小化输入/权重的量化误差
   把输入和权重拉伸成向量，量化得到，反量化后得到，也即
   
   
   其中，和分别是线性均匀分布的定点量化中的两个量化参数，分别用于缩放和偏移；
   为了方便讨论，我们令，那么
   
   
   那么可以定义一个误差函数，优化和校准过程通常就是寻找一个合适的量化参数，使得误差最小化。

2. 最小化输出误差
   记输入和权重为和，量化后为和，反量化后为和，卷积输出为；
   直观上看，最小化输出误差相比于直接最小化一个向量的量化误差更加直接，也更加有效，但实践上会发现可能并非如此，有时候会出现精度反而下降的情况，这可能是对小样本校准数据集过拟合的结果。此外，为了得到，优化和校准过程中每一次迭代都需要一次完整的前向传播（具体次数取决于迭代方法的收敛速度，或搜索方法的粒度），这意味着该方法的计算代价将远高于直接前一种方法。
   我们可以直接将优化目标定为最小化输出误差，也即最小化，这里的可以是

   1. ：保持全精度输入，只搜索确定的量化参数
   2. ：保持全精度权重，只搜索确定的量化参数
   3. ：同时搜索确定和的量化参数；或量化输入/权重其中一个，搜索确定另一个的量化参数
      

优化方法

1. 交替优化（Alternating Optimization）
   参考论文：《Fixed-point feedforward deep neural network design using weights +1, 0, and −1 (2014)》

假定其他参数确定且与当前参数无关，推导出待优化参数在当前情况下的最优解，然后固定该参数，继续推导下一个参数的最优解，反复迭代、交替优化，直到收敛为止（参数本身收敛，也即更新幅度很小；或者目标收敛，比如误差不再减小）。本质上是一种贪心搜索，由于实际场景中误差函数不是理想、光滑的凸函数，所以通常得不到全局最优解。

该方法需要求解迭代更新的公式，在前述的三种误差函数中，L2的迭代公式推导上最为方便，以下便用L2误差进行举例——（如果迭代更新公式难以推导，也可采用梯度下降法）

1. 以最小化向量的量化误差为例

   1. 初始化
      通常用朴素的最大最小值确定初始的
   2. 假设确定，并对向量进行量化
      
   3. 假设确定，推导出能使最小化的
      
      
      （这里假定了与无关；而且因为是个二次函数，极值点即为最值点）
      即为迭代更新公式
   4. 重复步骤ii和iii，直到收敛

2. 以最小化输出误差为例（）
   （此处假设没有偏置，如果有偏置，可以先把偏置减去之后优化过程同理）

   1. 初始化
      通常用朴素的最大最小值确定初始的
   2. 假设确定，并对权重向量进行量化
      
      
   3. 假设确定，那么也确定，推导出能使最小化的
      
      同理，迭代更新公式为
   4. 重复步骤ii和iii，但由于与有关，通常该优化过程是不稳定的，很难让收敛到很小的波动范围内。在实际应用中，更适合设置一个容忍性的超参patience，当连续patience次迭代都没有取得更小的，那么就停止优化过程
3. 黄金分割搜索（Golden Section Search）
   参考：《黄金分割搜索》
   黄金分割搜索是一种在单峰函数上搜索最值点的方法，搜索速度快，但非凸函数无法得到最优解。

已知一个单峰函数，且最小值点在区间内，为了搜索最小值点，我们需要在区间内引入第三个点

1. 如果，那么最小值点必定在区间内
2. 如果，那么最小值点必定在区间内

3. 如果且，那么最小值点可能在区间内，也可能在内。
   （假设在区间内取点）

   1. 如果，那么最小值点必定在区间内
   2. 如果，那么最小值点必定在区间内

重复该过程，即可逐步收缩区间，直到满足精度要求。


为了提高搜索效率，

1. 在较大的区间里选取点。
   如果，那么在区间上取点；否则在区间上取
2. 无论和大小关系如何，下一步的搜索区间是相等的。
   也即，那么可以得到点的确定方式为
3. 每次分割的区间应当成相等比例。
   比如（最小值点在区间内），要求，
   那么可以得到点的确定方式为为黄金分割比例

1. 栅格搜索（Grid Search）
   将搜索区间划分为若干栅格（通常是等间距的），依次计算误差函数，搜索出最佳的结果。该方法比较粗暴，搜索效果取决于栅格的密度，也即分辨率。栅格设置的越密集，效果越好但计算代价也越高，但好消息是，由于每个栅格是无关的，所以通常比较容易实现并行计算以加速优化过程。

数据收集方法

为了搜索输入的量化参数，我们往往需要若干不需要标注的样本来作为校准数据集，数据集的数量跟采用的方法有关，比如L2和cos的误差函数所需的数量就比较少，KL误差函数所需的样本数量就比较多。我们可以选择直接收集完整、准确的原始数据，也可以选择保留一些近似的数据，比如利用直方图的方法。

原始数据

收集原始数据是最直接的想法，但往往有些场合迫使我们无法存下如此之多的数据量，这主要原因是来自于内存（显存）的限制，毕竟保留大量样本在推理过程中的完整数据是需要大量内存开销的。


对于最小化L2距离和最大化余弦距离的方法，保留原始数据的动机很直观，那么对于本身就基于直方图的最小化KL距离方法有必要保留原始数据吗？其实，对于最小化KL距离方法，保留原始数据也有积极作用的，我们可以在每次确定量化的数值表示范围之后重新统计直方图来获取较高的分辨率（而直方图的方法是在现有直方图的基础上，直接去除一些bin，本质上又是一次近似过程）。

直方图

直方图的方法只需要记录每个bin上的数量，与校准集里的样本数无关，因此允许数量非常大的校准样本。当然，在使用直方图的时候也会遇到一些问题：

1. 不同mini-batch的直方图表示范围不同。
   对于不同的输入样本，卷积、全连接层的输入数值范围是不同的，这会导致对不同的mini-batch统计出来的直方图表示范围不同。为了将不同mini-batch的数据统一到一个直方图上，我们可以

   1. 把第一个mini-batch的数值表示范围作为最终直方图的表示范围，后续数据截断后直接放入直方图。也可以预留一定裕量，比如第一个mini-batch的数值范围为，那么可以将最终直方图的表示范围固定为
   2. 合并直方图，假设直方图中每个bin内是均匀分布的，然后将其分配到新的直方图上的各个bin中；在python上，这种逐个标量判断和处理效率是非常低的，如果使用了torch、numpy等能够批量数据的库，将原直方图采样若干倍再丢入新的直方图中是更加高效的选择（具体过程参考PyTorch-HistogramObserver的 _combine_histograms 方法）
2. 如何在直方图上计算L2误差？
   KL误差的计算过程本身就需要依赖于直方图，所以在直方图上计算KL损失的过程是非常直观的。而L2误差要如何计算？（具体实现可以参考PyTorch-HistogramObserver的 _non_linear_param_search 方法）
   

一些细节

量化顺序

对于校准的方法，包括输入量化和基于最小化输出误差的权重量化等需要样本数据进行校准的量化方式，需要考虑量化顺序的问题。

1. 数据在前向传播的过程中，权重是否要预先量化？
   通常来说，先量化权重，再校准输入是比较好的选择，这样在校准输入的过程中考虑了权重量化的影响，更符合最终量化的使用场景

2. 数据在收集的过程中，浅层特征图是否要进行量化？
   这个问题主要出现在基于直方图的数据收集方法中，由于数据需要分多个mini-batch输入，经过若干次数据收集后才能取得特征图的量化参数，因此直观的实现方式是在收集数据的过程中不考虑浅层特征图的量化。但通常而言，考虑浅层的量化影响才更符合最终量化的使用场景，

   1. 可以逐层进行量化，从浅至深每次只量化一层，这样一来在量化某一层时就已经确保浅层已经取得校准后的量化参数，能够进行正确的量化操作。但该方法在实际实现上会比较麻烦；
   2. 折中点的做法是，在数据收集过程中采用在线量化（即根据当前特征图的最大-最小值来确定量化参数），来引入量化的影响，在较高比特（如8bits）的量化方案中即使不对数据进行截断通常也不会有明显的负面影响，该方法一般也能工作的很好；
   3. 另一种折中的做法是，将数据校准过程分为两步，第一步先不考虑浅层量化的影响，通过校准得到一个相对准确的量化参数；第二步引入浅层量化的影响，用第一步取得的量化参数进行量化，进而对量化参数进行微调。实际应用中，也可以进行多次迭代，实现多轮的微调以期得到更优的结果，当然也可以结合上一种方法，在第一步校准过程中采用在线量化的方法

修正BN层

由于量化的影响，会导致特征图的数据分布发生变化（均值和标准差发生变化），这与训练时BN层统计的结果出现不一致，这尤其在低比特量化中将会有明显的负面影响。此时可以给BN层设置一个较大的momentum（甚至设为1），然后喂入一定数量的数据对BN层的统计量进行修正。要注意的是，用于修正BN层统计量的数据应当从训练数据集中随机抽取，从单一类别中抽取数据不具有代表性，很可能会导致统计量出现明显的偏差，修正后的模型表现反而比不修正的表现更差。

逐通道量化输入

逐通道量化（每个通道独享一套量化参数）的效果往往比逐层量化（每一层的所有通道共享一套量化参数）的效果更好，权重也经常采用逐通道的量化方式。《Low-bit Quantization of Neural Networks for Efficient Inference (ICCV2019)》在分类网络上做过一些int4量化权重的比较实验——



那么为什么大家倾向于逐通道量化权重，却对输入采用逐层量化呢？这主要是因为逐通道量化输入，在推理上不太友好——


首先，考虑逐层量化输入，逐通道量化权重的方案：


其中，是权重滤波器（也即输出通道）的索引号；


实际推理中，反量化时可以直接乘以预先计算好的，而不需要分开做两次乘法运算。


再看看逐通道量化输入的方案：


其中，是权重滤波器（也即输出通道）的索引号，是权重核（也即输入通道）的索引号；


主要问题在于那个，原本可以输出通道计算完之后再做反量化，但由于输入也做了逐通道量化，每算完一次输入通道都要先乘以来完成一部分反量化工作，然后才能求和得到输出通道，这是不太友好的。

但也不是说逐通道量化输入就不可行，《Towards Accurate Post-training Network Quantization via Bit-Split and Stitching (ICML2020)》就指出，我们可以把折叠到里（以下我们倒过来推导）

那么，

换句话说，就是把乘到权重的输入维度上去—— w[:, j, :, :] *= beta[j] 

微调权重

本文开头提到过，也有的后训练量化算法试图在优化时对权重的原始数值进行微调，这里简单介绍一下其中一种，来自中科院自动化所的Bit-Split&Stitch算法。

Bit-Split&Stitch

论文：《Towards Accurate Post-training Network Quantization via Bit-Split and Stitching (ICML2020)》
代码：https://github.com/wps712/BitSplit


简单来说，该算法采用交替迭代、最小化输出的L2误差的方式一边优化量化参数、一边微调权重。
但是！论文里的迭代更新公式推错啦！！！而且开源代码极其糟糕，都快读吐了。

整体设计

1. 比特分割
   记原始标量为，将其分割成；
   其中，为比特数，是有符号数；
   初始化时，取相同的符号
2. 逐比特优化
   输入校准样本，以最小化输出L2误差为目标，逐比特对权重进行优化
3. 比特粘合
   将逐比特优化好的合并回bits的整型标量

逐比特优化

首先定义一些符号，按照im2col对卷积权重和输入特征图进行展开，记输入特征图为，由于权重是逐通道（输出通道）量化的，我们不妨只分析其中一个通道的优化过程，也即，这并不影响我们推广到一般情况。量化后的权重记为；为简化推导过程，不妨假设输出特征图分辨率不变，记为，那么

先从优化目标入手，即希望量化权重之后能够最小化输出的L2误差

将量化后的逐比特分割，

，其中
优化目标展开为



的迭代公式很容易就可以推导出来

关键我们看看的迭代公式。这里并不好优化，Bit-Split&Stitch采用贪心搜索的方法，先假设其他比特都已经固定，然后单独优化这一比特，反复迭代直到收敛，每个子优化目标变为

其中，
，是量化因子和比特位所表示的数值的乘积；
，是输出对于权重第比特位的响应；
那么这个L2误差的子项为

注意到，，那么，而显然，
为了最小化，那么有

要注意到，推导出来的最终公式里包含了若干个项之和，这里实际上是假设其他项确定的前提下优化某一个子项，也即优化其中的某一元素（不同滤波器是独立无关的，可以并行优化）。

可能存在的问题

1. 是否过分假设？
   该方法不仅逐元素优化（假设其他元素确定的前提下优化单一元素），还进行逐比特优化（假设其他比特确定的前提下优化单一比特），假设性非常强，这波贪心搜索非常贪心
2. 涉及权重本身微调，是否容易出现过拟合？
   校准集通常挑的不多，如果有小样本进行校准来微调权重，可能很容易出现过拟合；如果用大样本来微调权重，那为什么不直接做重训练呢
3. 优化过程非常缓慢
   该优化过程非常复杂，所需时间也很长，远高于常见的post-training方法，开源的代码更是糟糕——一个ResNet18的校准过程我在V100上整整跑了三四个小时？！