模型量化-学习笔记

一. 基本概述

1.1 背景

随着深度学习的不断发展，神经网络广泛应用于不同的领域，取得远超以往的效果的同时深度网络模型的复杂度也越来越高, 这严重制约了它在工业界发展.

模型量化是目前工业界最有效的模型优化方法之一.

1.2 定义

模型量化是指把模型的参数从FP32映射到nbit位的过程, 简单来说就是在定点数与浮点数等数据之间建立一种数据映射关系, 使得以较小的精度损失代价获得了较好的收益。 例如FP32-->INT8可以实现4倍的参数压缩，在压缩内存的同时可以实现更快速的计算，从而有效地提高模型的性能; 最极端的二值量化理论上甚至可以实现32倍的压缩，但是过度的压缩会导致模型的精度快速下降, 所以更多量化的时候需要做好精度和性能的权衡.

工业界一般使用int8量化, 在模型推理前需要把FP32映射为int8进行计算, 然后在输出的时候做一个去量化操作, 把计算的int8结果映射回FP32.

image.png

在介绍量化原理之前, 我们至少需要回答3个问题:

• 为什么不直接训练一个低精度的网络? 一般训练使用的是梯度下降法, 直接用 定点数做参数训练很容易导致梯度消失, 训练的模型会欠拟合, 所以训练阶段我们还是要用 浮点数 来做参数.
• 为什么可以做量化? 神经网络模型一般有较好的抗干扰能力对噪声不敏感，量化相当于对原输入加入了大量的噪声, 对模型的精度一般不会造成太大影响.
• 为什么用低精度数量化后性能可以提升? 首先我们要知道一点就是: 模型的性能主要由模型的参数量和计算量来决定的. 参数量的大小直接决定模型的大小，也影响推断时对内存的占用量以及内存的访问次数. 很多时候模型计算需要内存和CPU共同协作, 同一个计算量下, 模型越大内存占用就越大内存访问次数也就越多, 而我们知道CPU和内存的速度根本就不在一个数量级, 所以模型的参数量大小一般决定了模型性能的下限.

计算量越大的模型, 性能一般越差, 我们知道CPU或者GPU的FLOPS(每秒浮点计算量)是相对固定的, 模型的计算量大小是决定了模型性能的上限.

• 要量化的模型参数主要指哪些参数? 量化的参数主要分为2类: 权重和激活值, 其定义如下图所示:

image.png

1.3 作用

• 减少内存带宽和存储空间
• 提高系统吞吐量（throughput），降低系统延时（latency）
• 易于在线升级，模型更小意味着更加容易传输
• 减少设备功耗，内存耗用少了推理速度快了自然减少了设备功耗
• 支持微处理器，有些微处理器属于8位的，低功耗运行浮点运算速度慢，需要进行8bit量化

1.4 挑战

• 定点数/FP16表示范围和精度完全不同明显比FP32小, 如果使用FP16的话, 还需要硬件级别的支持.

image.png

通过上图, 我们可以看到不同的数据类型的表示范围和精度完全不同, 简单地做类型转换显然是不行的.

• 如何保证量化后的精度损失是最低的?

1.5 解决方案

• 根据模型不同的参数的统计情况, 选择不同的参数量化方案.
• 最大化模型量化前后参数分布的相似性, 从而将量化的损失最小化

二. 量化的分类

量化的本质就是对float32进行缩放, 即: Q=R/S+Z

R 表示真实的浮点值，Q 表示量化后的定点值，Z 表示 0 浮点值对应的量化定点值，S 则为缩放因子。

2.1 均匀量化/非均匀量化

按照量化级的划分方式来分, 假设现在要把FP32映射到int8, 那么有:

• 均匀量化: 动态范围被均匀地划分为128份, 其中128表示量化级。
• 非均匀量化: 动态范围的划分不均匀, 一般用类似指数的曲线进行量化或者使用Kmeans对网络权重和特征进行聚类，得到不同的聚类中心， 然后将聚类中心作为同一簇权重的量化代表。

image.png

2.2 对此量化/非对称量化

根据Z(零点)的不同, 我们又分为对称量化和非对称量化

• symmetric quantization(对称量化): FP32的0点被映射到了int的0点
• asymmetric quantization(非对称量化): FP32的0点没有映射到int的0点

image.png

2.3 饱和量化/非饱和量化

按照阈值选择的不同划分, 我们又分为饱和量化/非饱和量化

image.png

• 非饱和量化 ( No Saturation) : 计算FP32类型Tensor中绝对值的最大值abs_max，将其映射为127，则量化比例因子等于abs_max/127。 (1) 对于权重参数一般分布比较均匀, 所以采用非饱和量化方法。 (2) 当正负分布不均匀的时候，是有一部分是空缺的，也就是一部分值域被浪费了, 所以说这种量化方式是不饱和的。
• 饱和量化 (Saturation): 使用通过统计计算一个合适的阈值T (0<T<mab_max)，将其映射为127，则量化比例因子等于T/127。 对于激活值Tensor（包括输入和输出）其值分布不均匀，所以采用饱和量化(最大值量化)方法 。

2.4 训练后量化PTQ

(1) 定义

在模型训练好后量化, 量化其实就是权重和激活值进行缩放的过程，在PTQ中, 我们是通过统计的方法, 使用Calibration Dataset来近似模拟现实的数据分布, 从而得到

权重和激活值的动态范围(Gather layer statistics) 和量化参数（q-parms), 以此来对我们的权重和激活值进行合理的缩放。

(2) 量化的过程

• 使用校准数据集, 这个数据集要经可能具有多样性, 有代表性, 理想情况下是验证数据集的子集, 对已经预训练好的模型的每一个layer进行统计。
• 基于上一步的统计来 权重和激活值 的动态范围(Gather layer statistics) 和量化参数（q-parms）
• 使用确定好的动态范围(Gather layer statistics) 和量化参数（q-parms）对模型的 权重和激活值 进行量化。

image.png

(3) PTQ的特点

• 速度快, 无需重新训练模型。
• 量化方案的即插即用, 方便进行组件化。
• 对量化后模型精度控制较弱。 在某些位置, 为了保证精度我们需要保持用FP32计算, 有些位置为了速度直接用int8计算就行, 这个在PTQ中无法实现. 理想的状态下应该是可以指定量化和去量化发生的位置, 同时又能够让模型主动学习到权重和激活值的缩放参数。

image.png

2.5 训练中量化(QAT)

(1) 定义

在模型训练阶段进行量化, 由于PTQ可能存在一些误差，所以我们需要一种可学习的scale。

训练中量化QAT 就是在做这样一件事情。简单概括就是，我们在网络训练过程去模拟量化，我们通过设定一个可学习的scale，这个scale一般可以与weights或者激活值相绑定，然后我们利用一个量化过程 q = round(r/s)127，将需要量化的值量化到0, 127之间，再接着一个反量化过程q s，就实现了一个误差的传递，接着我们利用反量化后的结果继续前传，最后得到loss，我们求量化后权重的梯度，并用它来更新量化前的权重，使得这种误差被网络抹平，让网络越来越像量化后的权重靠近，最后我们得到了量化后的权重q和缩放因子s。而这一系列操作都可以写成网络中的一个op，实现网络的正常训练。最后，我们利用q和s，来进行线上推理。

(2) 量化过程

• 从预训练模型开始，在不同网络层中添加量化操作
• 利用若干epoch模型进行调优, 模拟在推理过程中发生的量化过程
• 通过训练学习量化参数，减少量化模型和与预训练模型之间的精度损失.

image.png

(3) QAT的特点

• 速度慢, 需要重新训练模型。
• 方案的即插即用, 方便进行组件化。(需要重新训练)
• 对量化后模型精度控制较好, 直接对量化后的损失进行建模, 理论可以获得最好的速度和最小的精度损失。

QAT中我们可以直接对指定位置的指定操作进行量化或者去量化, 这完全取决于我们是希望获得精度还是速度.

image.png

三. TensorRT/NIC量化方案(int8量化)

3.1 权重值量化

image.png

如上图所示, 一个预训练好的模型, 它的权重一般分布比较聚集, 比较均匀, 所以这部分Nvidia直接使用不饱和量化&对称量化, 对称量化的Z等于0, 所以只需要求scale, 具体操作如下:

image.png

据说刚开始Nvidia还使用了偏置项(如左图), 后来发现这一项对模型精度没有影响就去掉了(如右图), 其实就是做简单的线性缩放, 其中 $scale_A =\frac{1}{|max_A|}, $ $scale_B =\frac{1}{|max_B|}$, scale_A, scale_B表示的是缩放因子.

3.2 激活值量化

真正麻烦的是激活值的量化,Nvidia统计了3种模型的不同layer的激活值分布情况如下:

image.png

这里的x轴表示的是激活值, y轴是激活值的出现的次数次数统计后进行标准化的结果, 每一个样本的样本参数分布都是一个散点图, 由于神经网络一般使用Relu做中间层的激活函数, 通过下面的函数图可以看到这个函数的激活值有下界但是没有上界, 所以很多时候激活值的范围都是比较大的.

image.png

事实也是如此, 通过上面的3个激活值分布图可以看出:

• 激活值的分布范围一般都比较广, , 这种情况下如果直接使用不饱和量化的话, 就会把离散点噪声给放大从而影响模型的精度.
• 不同模型的不同层的分布差异也非常大, 所以需要对每个模型的每一层都有一个阈值这种量化方式叫做逐层量化, 也可以对每一层每个通道都进行独立量化，这对精度也会有一个很好的提升。

由于存在离散点噪声的原因, Nvidia在对激活值量化时, 不再是基于最大值进行缩放, 而是找一个阈值|T|, 以此来控制量化的float范围, 避免离散点噪声被放大,

通过前面的介绍我们也知道这种方法叫做饱和量化, 关键是如何找到最优的阈值?

image.png

TensorRT 为了找到一个最优的阈值, 引入了KL散度(也称相对熵)来衡量不同的INT8分布与原来的FP32分布之间的差异程度, 不了解这个熵理论的可以看这里信息熵理论, 它的公式如下:

KL_{divergence}(P,Q):= \sum_{i=1}^n (P[i]*log(P[i]/Q[i])

其中P，Q分别称为实际分布(reference_distribution)、量化分布(quantize_distribution), KL散度越小, 说明两个分布约接近, nvidia把确定每一层的|T|值的过程被称为校准(Calibration), 大体的流程如下:

image.png

• 从验证集选取一个子集作为校准集(Calibration Dataset )，校准集应该具有代表性，多样性，最好是验证集的一个子集。
• 把校准集输入到模型进行前向推理, 并收集模型中各个Layer的激活值分布直方图, 如左上图所示.
• 然后生成不同的阈值和量化分布, 并计算不同阈值下真实分布P和量化分布的Q的KL散度, 直到找到KL散度最小的一个阈值.

TensorRT官方也给出了这一个过程的伪代码:

image.png

nvidia给出了一个demo实例:

import scipy.stats as stats
import numpy as np
test_p = [1, 0, 2, 3, 5, 3, 1,  7]
spilt_array = np.array_split(test_p, 2)
ret_q = []
for arr in spilt_array:
    avg = np.sum(arr)/np.count_nonzero(arr)
    for item in arr:
        if item!=0:
            ret_q.append(avg)
            continue
        ret_q.append(0)
print(ret_q) 
test_p/=np.sum(test_p)
ret_q/=np.sum(ret_q)
print("KL(P,Q)=", stats.entropy(test_p, ret_q))

image.png

具体细节如下:

step1: 首先将激活值分成2048组, 每组包含多个数值, 注意如果激活值的范围大小没有超过128, 那么是可以直接进行映射的.

step2: 不断地截断参考样本P，长度从128开始到2048, 截断区外的值加到截断样本P的最后一个值之上；最后求得样本P的概率分布；

step3: 创建样本分布Q，其元素的值为截断样本P的int8量化值, 将Q样本长度拓展到 i ，使得和原样本P具有相同长度；求得Q的概率分布 并计算P、Q的KL散度值

step4: 不断循环step2, step3, 就能不断地构造P和Q并计算相对熵，然后找到最小（截断长度为m）的相对熵, 而阀值就等于（m + 0.5）*一个bin的长度；

校准后的效果:

image.png

更多关于KL散度的说明请参考附录.

3.3 量化后的最终效果

• 模型精度上, 量化后对模型的精度影响比较小, 甚至可能会带来精度的提升.
• 模型性能上, 量化后一般随着batch_size提高一般会有1.5~3.5倍的性能提升.

image.png

3.4 存在的问题

• TensorRT的KL量化方案只是针对了Relu这种激活函数, 这种情况下激活值的取值在0到正无穷. 刚好落在正半轴, 实际场景要考虑激活值在负半轴的情况. 而且左右两侧的边界分布还不一样.

image.png

• TensorRT的KL量化时没有考虑KL散度平滑问题, 根据KL散度的公式:
  

  image.png
  在边界问题, 当p(x)>0, p(y)=0时, H(X, Y) 会无穷大，这就好比一个分布(X)认为某个事件可能发生只是概率值小，但是另外一个分布(Y)却认为该事件不可能发生，因此这两 个分布是语义上时完全互斥的, KL大也很正常. 但是要知道X,Y是参考分布和量化分布, 他们本来就是来自同一个分布的概率分布和频率分布, 出现这种情况说明量化的时候 y 这个点的没有采到样本(只是没采集到样本, 和量化的区间有关, 有可能被分到隔壁区间去了), 为了纠正这种误差我们就需要对样本的0点进行平滑处理, 或者直接抛出异常.

可以参考一下neural-compressor里面的实现: https://github.com/intel/neural-compressor/blob/master/neural_compressor/utils/kl_divergence.py

四. 附录

1. KL散度的实例1

KL分布是用来衡量2个数据分布的距离的, 当两个分布越相近, KL值越小, 分布一样的话, KL=0, 其公式为:

H(f|g) = D_{KL}(f, g) = \sum_{x\in X}f(x)*log_2\ \frac{f(x)}{g(x)}

假设f(x)为真实分布X，g(y)为非真实分布Y，p表示事件发生的概率, 则上述相对熵为：

H(X|Y) = D_{KL}(X||Y) = \sum_{i=1}^{N}\ p(x)\ log_2\ \frac{p(x)}{p(y)}

所以我们可以用KL散度来找最优的分布情况, 例如现在我们随机一组数据表示原始数据的分布情况

import numpy as np
SIZE = 10
x = [np.random.uniform(1, SIZE+1) for i in range(SIZE)]
px = x / np.sum(x)
def get_KL():
    y = [np.random.uniform(1, SIZE+1) for i in range(SIZE)]
    py = y / np.sum(y)
    
    KL = 0
    for i in range(SIZE):
        KL+=px[i]*np.log(px[i]/py[i])
    return KL, y

然后我们可以使用暴力法找到一个KL在可接受范围内的分布:

from itertools import count
import matplotlib.pyplot as plt
def get_KL_Q(kl_threshold=0.01):
    for i in count():
        KL, y = get_KL()
        if KL <  kl_threshold:
            plt.plot(x)
            plt.plot(y)
            break
    return

我们看一下 kl_threshold=0.01时, y的分布情况:

get_KL_Q(0.01)

image.png

我们看一下 kl_threshold=0.05时, y的分布情况:

get_KL_Q(0.05)

image.png

显然, KL越小, Y和X就拟合的越好.

2. KL散度的示例2

假设现在有一个float32的激活值的分布如下:

step1. 生成数据分布P, 并绘制P的概率分布直方图

import random
import numpy as np
import matplotlib.pyplot as plt 
def generator_P(size):
    walk = []
    avg = random.uniform(3.000, 600.999)
    std = random.uniform(500.000, 1024.959)
    for _ in range(size):
        walk.append(random.gauss(avg, std)) 
    return walk

size = 20480
P = generator_P(size)
P = np.array(P)
P = P[P>0]
print("最大的激活值", max(np.absolute(P)))
 
# 可视化
plt.title("Relu activation value Histogram")
plt.xlabel("Activation values")
plt.ylabel("Normalized number of Counts")
# bins表示要分成多少组, density表示是否要normed
n, bins, _ = plt.hist(P, bins=2047, density=True)
plt.show()
print("各个组的范围", bins)
print("各个组包含的样本概率", n)
print("步长:", bins[2]-bins[1])

image.png

step2. 解决KL散度量化的平滑问题

def smooth_distribution(p, eps=0.0001):
    """
    给定离散分布（可能尚未标准化为1），通过将0替换为eps乘以比例因子使其平滑, 并从非零值中减去相应的量。
    """
    is_zeros = (p == 0).astype(np.float32)
    is_nonzeros = (p != 0).astype(np.float32)
    n_zeros = is_zeros.sum()
    n_nonzeros = p.size - n_zeros
    if not n_nonzeros:
        raise ValueError('The discrete probability distribution is malformed. All entries are 0.')
    eps1 = eps * float(n_zeros) / float(n_nonzeros)
    assert eps1 < 1.0, 'n_zeros=%d, n_nonzeros=%d, eps1=%f' % (n_zeros, n_nonzeros, eps1)
    hist = p.astype(np.float32)
    hist += eps * is_zeros + (-eps1) * is_nonzeros
    assert (hist <= 0).sum() == 0
    return hist

import matplotlib.pyplot as plt
import copy

test_p = np.array([0, 2, 3, 4, 2.4 ,5, 6])
plt.plot(test_p, color='r')

smooth_test_p = copy.deepcopy(test_p)
smooth_test_p = smooth_distribution(temp_p, eps = 0.5)
plt.plot(smooth_test_p, color='g')
plt.show()

step3. 使用KL散度量化求阈值

import copy
import scipy.stats as stats
def threshold_distribution(distribution, target_bin=128):
    distribution = distribution[1:]
    length = distribution.size
    threshold_sum = sum(distribution[target_bin:])
    kl_divergence = np.zeros(length - target_bin)
    
    for threshold in range(target_bin, length):
        sliced_nd_hist = copy.deepcopy(distribution[:threshold])

        # generate reference distribution p
        p = sliced_nd_hist.copy()
        p[threshold - 1] += threshold_sum
        threshold_sum = threshold_sum - distribution[threshold]

        # is_nonzeros[k] indicates whether hist[k] is nonzero
        is_nonzeros = (p != 0).astype(np.int64)
        
        quantized_bins = np.zeros(target_bin, dtype=np.int64)
        # calculate how many bins should be merged to generate 
        # quantized distribution q
        num_merged_bins = sliced_nd_hist.size // target_bin

        # merge hist into num_quantized_bins bins
        for j in range(target_bin):
            start = j * num_merged_bins
            stop = start + num_merged_bins
            quantized_bins[j] = sliced_nd_hist[start:stop].sum()
        quantized_bins[-1] += sliced_nd_hist[target_bin * num_merged_bins:].sum()

        # expand quantized_bins into p.size bins
        q = np.zeros(sliced_nd_hist.size, dtype=np.float64)
        for j in range(target_bin):
            start = j * num_merged_bins
            if j == target_bin - 1:
                stop = -1
            else:
                stop = start + num_merged_bins
            norm = is_nonzeros[start:stop].sum()
            if norm != 0:
                q[start:stop] = float(quantized_bins[j]) / float(norm)
        
        # 平滑处理, 保证KLD计算出来不会无限大
        p = smooth_distribution(p)
        q = smooth_distribution(q)

        # calculate kl_divergence between q and p
        kl_divergence[threshold - target_bin] = stats.entropy(p, q)

    min_kl_divergence = np.argmin(kl_divergence)
    threshold_value = min_kl_divergence + target_bin

    return threshold_value

step4. 阈值可视化

# 获取KL最小的阈值
hist, bins = np.histogram(P, bins =2048)  
threshold = threshold_distribution(hist, target_bin=128)
print("threshold 所在组:", threshold)
print("threshold 所在组的区间范围:", bins[threshold])
# 分成split_zie组, density表示是否要normed
plt.title("Relu activation value Histogram")
plt.xlabel("Activation values")
plt.ylabel("Normalized number of Counts")
plt.hist(P, bins=2047)
plt.vlines(bins[threshold], 0, 30, colors = "r", linestyles = "dashed")
plt.show()

image.png

五. 参考

《TensorRT 量化实现细节》

《TensorRT INT8量化》

《Google量化白皮书》

《mxnet量化实现》