量化对比

根据训练过程的不同,量化可分为从头训练(train from scratch)、重训练(retrain)、后训练(post-training)三种。
根据量化方法的不同,又可分为非线性量化线性量化

非线性量化:《Deep Compression - Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding》这一经典的模型压缩开山鼻祖的论文中使用的K-Means聚类方法就是典型的非线性量化,不过该种非线性量化依赖聚类中心的查表,只能起到损失尽可能少的精度而压缩模型大小的效果,并不能对运算进行加速。
线性量化:模型量化 - 图1,以线性间隔对向量进行量化,可以将卷积操作从fp32浮点运算转为int8/uint8甚至int4/uint4等的整形运算,不但能压缩模型大小,还能在嵌入式设备上节省能耗并对网络推理进行加速。

从头训练:暂时没了解到这种方法有什么应用场景。
重训练量化:在训练好的float模型上引入模拟量化,对网络进行重训练,同时调整各层的scale,一般该恢复训练过程只需要非常少的epoch数。重训练量化的缺陷是需要整个训练数据集与充足的算力
后训练量化:对训练好的float模型使用合适的量化操作和校准方法,以实现量化损失最小化。该过程不需要训练,通常不需要更新权重的原始数值而只是为了获取合适的量化scale。后训练量化的优势是只需要非常少的输入数据用于校准,不依赖于大型计算框架,不要求高算力;缺陷是准确度相对不如重训练量化

重训练量化

重训练量化,QAT(Quantization-Awared Training)最常见的形式,具体参见04-模型量化里面记录了自己瞎写的一个QAT小项目的历程和实验结果。

后训练量化

引用郑老师博文:后训练量化
里面很多方面都讲的更加详细,这里主要收录了重点。

误差函数

误差函数的度量主要包括:L2距离、KL散度、余弦距离。它们都是用于度量量化前后的张量间的差距,显然这个差距越小越好,我们的量化过程也需要使得这个loss尽可能小。

L2距离:相对直观,也最常用,易于求导与优化,数学性质良好。L2距离也有其它变种,譬如从裁剪的经验上来看我们往往认为更大的权重更重要,因此可以在计算L2 loss时对权重大小进行加权,该变种在低比特量化上似乎能带来更好的结果,高比特时反而会适得其反(没做实验,存疑)。
KL散度:对数据分布非常敏感,需要较大数据集进行校准。根据经验值,用于L2距离校准只需要几十张图,但用KL散度就需要一千张。该误差函数由TensorRT推动,具体见基于KL散度的离线量化
余弦距离:没具体了解过,和本科毕设推荐算法里用的度量方法一样。余弦距离的取值范围为[-1, 1],是越大越好而非越小越好。

优化对象

最小化量化误差
直接以量化前后的激活值、权重的距离作为误差函数,对scale进行优化。对于权重的优化,直接离线原地优化即可;对于激活值,可以喂入少量数据用于校准。

最小化输出误差
记输入的激活值为模型量化 - 图2,权重为模型量化 - 图3;量化后为模型量化 - 图4模型量化 - 图5,全精度的卷积的输出为模型量化 - 图6,而经过量化的卷积输出为模型量化 - 图7。在最小化输出误差时,我们的优化目标为最小化模型量化 - 图8模型量化 - 图9的误差。其优化过程可以为同时优化模型量化 - 图10模型量化 - 图11,也可以交替地保持一方为全精度输入,优化另一方。
直观上,最小化输出误差显得更为有效,但实践上可能并非如此,这可能是对于用于校准地小样本过拟合的结果。

优化方法

梯度下降
梯度下降是最简单粗暴的一种方法,但是需要用非常小的学习率而且非常不稳定,需要进行一定的监督。其优势是适用于所有的误差函数,但是依赖于大规模计算框架的自动求导。(郑老师的说法是,做后训练量化的一大优势就是轻资源,很多客户不愿意把数据集与项目交给我们部署,他们需要的只是部署工具而不具备算力和计算框架)

  1. scale = Parameter(torch.Tensor([MAXMIN_OBSERVE_VALUE])) # 用传统max-min值得到的scale初始化
  2. optimizer = optim.SGD([scale], lr=3e-9, weight_decay=0., momentum=0.) # L2和动量必须为0
  3. for i in range(100):
  4. q_weight = layer.weight / scale
  5. q_weight = Round.apply(q_weight)
  6. q_weight = torch.clamp(q_weight, -8, 7)
  7. q_weight *= scale
  8. loss = ((q_weight - layer.weight) * (q_weight - layer.weight)).sum()
  9. loss.backward()
  10. optimizer.step()
  11. print(loss.item(), scale.item())


交替优化**
假定其它参数确定且与当前参数无关,推导出待优化参数在当前情况下的最优解,然后固定待优化参数获得其它参数,继续推导下一个待优化参数,直到收敛。该方法本质上是一种贪心搜索,因此通常无法得到全局最优解。
这里我们以L2距离为误差函数举例,假设模型量化 - 图12,误差函数为模型量化 - 图13,我们假定模型量化 - 图14模型量化 - 图15无关,推导出当前最优的模型量化 - 图16后重新计算模型量化 - 图17,直到收敛。该方法相对梯度下降法要稳定不少,但是要求误差函数有在局部条件下易于求取极值。
L2距离误差函数是一个二次函数,其极小值点可根据二次函数公式得出。但其它损失函数可能就没这么容易推导局部条件下的最优解,难以使用交替优化方法。

  1. scale = Parameter(torch.Tensor([MAXMIN_OBSERVE_VALUE])) # 用传统max-min值得到的scale初始化
  2. for i in range(200):
  3. q_weight = layer.weight / scale
  4. q_weight = Round.apply(q_weight)
  5. q_weight = torch.clamp(q_weight, -8, 7)
  6. scale[0] = (layer.weight * q_weight).sum().item() / (q_weight * q_weight).sum().item()
  7. q_weight *= scale
  8. loss = ((q_weight - layer.weight) * (q_weight - layer.weight)).sum()
  9. print(loss.item(), scale.item())

栅格搜索
就是暴力搜索,但是好在每个栅格都是相互独立的,易于并行计算。

数据收集

不同于权重的量化过程只需要简单地对固定地权重值确定scale,激活值的量化相对复杂,往往需要若干样本用于校准。用于校准地数据集规模和采用地误差函数有关,L2和余弦误差函数需要的校准集较小,KL则较大。
有的时候校准数据集过大,受限于内存(显存)大小,无法一次性投入网络用于确定各个量化scale——即使使用L2作为误差函数,对于一些目标检测任务,数十的batch-size依然无法被有限的显存存储。一种折中的方法是将校准集也分割为mini-batch,每个mini-batch前向时计算当前的最佳scale,不同mini-batch间做滑动平均,但是似乎没人这么做过。

直方图法
直方图法通过统计各层激活值的分布,使得数据收集操作随校准集大小的空间复杂度从模型量化 - 图18下降为模型量化 - 图19,可以将非常大量的校准集信息融合在有限的空间中。
直方图法只需要记录整个校准集中的图片在前向过程中,各层特征图的值在bin中的分布情况即可。但是使用直方图法我们会碰到一些问题:

  1. 不同mini-batch的直方图范围(直方图的min、max)不同

    1. 以第一个mini-batch的范围确定最终直方图的范围,超出的全部截断。可以以第一个mini-batch的范围,上下留出一定余量,最终直方图的范围确定为模型量化 - 图20
    2. 对于一个新的mini-batch,检测当前batch的min、max值是否位于原有直方图范围内,若位于原有范围内则直接相加;对于超出原有范围的情况,对直方图重新采样后合并为当前mini-batch范围的直方图。Pytorch的HistogramObserver就采用了这种方法:参考以下的简化代码 ```python def combine_histogram(new_hist, orig_hist, down_rate假设=150, up_rate=128):

      原有直方图上采样为2048*128 bins

      upsample_hist = orig_hist.repeat_interleave(up_rate) # [1,2,3]->[1x128,2x128,3x128]

      新的直方图上采样为2048*150 bins(150这个数字是取上界取整得到)

      hist_with_output_range = torch.zeros(self.bins * down_rate)

      将原有直方图复制到新直方图的对应位置,对应位置由start_idx控制

      hist_with_output_range[start_idx : self.bins*up_rate+start_idx] = upsample_hist

      将原有直方图下采样,cumsum用于从a[3]获取[a[0],a[0]+a[1],a[0]+a[1]+a[2]]

      这一步将原有直方图在新的range下,从2048*150 bins通过每150个bins求和下采样为2048

      integral_hist = torch.cumsum(hist_with_output_range, dim=0)[down_rate-1 :: down_rate]

      错位做差

      shift_hist = torch.zeros(self.bins) shift_hist[1:] = integral_hist[:-1] interpolated_hist = (integral_hist - shifted_integral_hist) / upsample_rate

      return new_hist + interpolated_hist

def forward(x): cur_max, cur_min = x.max(), x.min if cur_min >= self.min and cur_max <= self.max: self.hist += torch.histogram(x, self.bins, min=self.min, max=self.max) else: cur_min, cur_max = min(cur_min, self.min), max(cur_max, self.max) down_rate = torch.ceil((cur_max - cur_min) / (self.max - self.min) * self.up_rate) self.min, self.max = cur_min, cur_max new_hist = torch.histogram(x, self.bins, min=self.min, max=self.max) self.hist = self.combine_histogram(new_hist, self.hist, down_rate, self.up_rate) ``` combine_histogram中都是在“将原来的直方图在新的range下重新采样”,相比于python中逐个bin加权判断效率低下,这种粗暴的先上采样再下采样的方法反而由于有算子优化而效率更高。pytorch对up_rate的取值为128,up_rate越小那么直方图需要添加的无效区间越多(因为down_rate依赖于up_rate且需要取整)。

  1. 在直方图上计算L2误差

引用郑老师的公式,如下。其中第四行的模型量化 - 图21为概率密度,由于我们假设bin内部是均匀分布的,因此可以用常数(bin的宽度的倒数)代替。以下公式只适用于整个bin都归于同一个模型量化 - 图22的情况!!!公式的第三行中,若碰到直方图中有bin,譬如左边一部分量化为3,但右边一部分量化为4,甚至有bin跨越了多个量化级数,需要分开讨论后将loss求和。
参考Pytorch的HistogramObserver_non_linear_param_search方法源代码。
模型量化 - 图23

量化细节

量化顺序

一般进行后训练量化时,先量化权重再校准输入激活值是较好的选择,除非以最小化输出误差为优化对象。但是我们在前向过程中进行量化时,深层的特征图没有引入浅层卷积的量化,但通常考虑浅层的量化才更符合最终量化的使用场景。

  1. 逐层量化,从浅至深每次只量化一层,量化某一深层时所有的浅层都已经经过了量化处理。这是最优的方法,但是实现较为复杂。
  2. 折中的做法,在数据收集过程中采用max-min的naive方案量化,以此引入量化的影响。该方法在较高比特(8bit)量化时一般工作良好。
  3. 另一种这种做法(感觉这种听着比较靠谱),首先不考虑浅层量化得到一个相对准确的校准,而后利用之前得到的量化scale再重新进行一次数据收集后进行校准。此过程可持续多次直到收敛。

修正BN层

由于量化的影响,特征图的数据分布可能发生变化,这与训练阶段BN层的统计结果可能不一致,这尤其在低比特量化时会引入非常明显的负面影响。此时可以给BN层设置一个非常大的momentum后,喂入校准数据重新对BN层的统计量进行修正。需要注意的是,用于修正BN层统计量的数据应从训练集中随机抽取(涵盖所有分类),单一类别的数据不具有代表性,偏置较大,可能导致明显的统计量偏差,反而使得修正后的模型劣化。