一文读懂TensorRT整数量化

接下来有空也会整理一些实战性的东西，比如结合pointpillars网络，用TensorRT进行PTQ int8量化和利用pytorch_quantization进行QAT量化。感兴趣可以关注下！

待继续整理

量化简介

模型量化是一种流行的深度学习优化方法，其中模型数据（包括网络参数和激活）从浮点表示转换为较低精度表示，通常使用 8 位整数。这有几个好处：

• 在处理 8 位整数数据时， NVIDIA GPU 使用更快更低成本的 8 位张量核来计算卷积和矩阵乘法运算。这会产生更多的计算吞吐量，这在计算受限的层上尤其有效。
• 将数据从内存移动到计算单元（在 NVIDIA GPU s 中的流式多处理器）需要时间和能量，而且还会产生热量。将激活和参数数据的精度从 32 位浮点值降低到 8 位整数可导致 4 倍的数据缩减，从而省电并减少产生的热量。
• 有些层有带宽限制（内存有限）。这意味着它们的实现将大部分时间用于读写数据，因此减少它们的计算时间并不会减少它们的总体运行时间。带宽限制层从减少的带宽需求中获益最大。
• 减少内存占用意味着模型需要更少的存储空间，参数更新更小，缓存利用率更高等等。

缺点：模型模型量化-反量化（quantization-dequantizaion）操作取整和截断会引入误差。

量化的两个重要过程，一个是量化（Quantize），另一个是反量化（Dequantize）：

• 量化就是将浮点型实数量化为整型数（FP32->INT8）
• 反量化就是将整型数转换为浮点型实数（INT8->FP32）

量化可分为两步：

• 选择要量化的实数范围，限制在该范围之外的值
• 将浮点型实数量化为整型数（比如FP32->INT8）

在预先训练的浮点神经网络中启用整数运算需要两个基本操作：

• 将实数转换为量化的整数表示（例如从 FP32 到 INT8）
• 将数字从量化的整数表示转换为实数（例如从 INT32 转换为 FP16）

量化类型

模型量化方法本质是函数映射，根据函数是否线性可以把量化分为线性量化和非线性量化

非线性量化与线性量化

“非线性”映射函数多种多样，通常需要根据不同场景的权值输入分布特点，研究使用何种映射方式。
非线性映射最显著的一个特点是，可以把不同重要程度的权值，映射到不同的量化范围。假如，权值输入主要分布在$[X_m，X_n]$中，其余区间的权值虽然比例较少，但也不能忽略，则可以通过非线性函数，把[，]区间中的权值映射到更大的量化后区间内，增加训练过程对于主要权值分布的敏感性。

在模型量化早期，会使用聚类方法，把权值聚为若干个类别，然后每个类别量化为同一个定值。这也一种非常典型的非线性量化方式。

理论上：非线性量化理论上精度更高（实际不一定），但是实现难度更大。实践中大多都采用线性量化方案。

线性量化是目前常用的方法，工业界应用比较成熟的8 bit量化方案采用都是线性量化

非对称量化和对称量化

是否以0为中心，0两边的动态范围是否一致。 $q=x.s+z$，对称量化$z=0$ ，图左边是非对称量化，右边是对称量化

训练后量化和量化感知训练

按量化是否参与训练可分为：训练后量化（PTQ）和量化感知训练（QAT）

• 训练后量化（PTQ），PTQ不需要训练，是一种轻量级的量化方法，在大多数情况下，PTQ足以实现接近FP32性能的INT8量化，针对更低位的量化又局限性；
• 训练时量化（QAT)，它可以获得高精度的低位量化，缺点：需要修改训练代码，并且反向传播对梯度的量化误差很大，也容易导致不能收敛。

饱和量化和非饱和量化

如下图左所示，简单的将一个 tensor 中的 和 FP32 类型的值映射为 -127 和 127 ，中间值按照线性关系进行映射，称这种映射关系为不饱和的。由于模型数据分布可能是不均匀的，可能存在一些离群点，其数值较大，但对结果影响又很小，直接使用非饱和量化会使得量化后的值都挤在一个很小的范围从而造成精度损失。

如上图右所示，TensorRT 的 INT8 量化采用对称饱和线性映射的量化方式

饱和线性映射方式不是将 $|\pm max|$映射为$|\pm 127|$，而是存在一个阈值 |T|（|T| < |max| ），将 |T|映射为$|\pm 127|$ ，超出$\pm |T|$阈值 外的直接映射为阈值$\pm 127$ 。

饱和量化会采用一些策略对数据进行截断，去掉了一些不重要的因素，保留了主要成分，减少精度的损失。比如TensorRT int8 校准使用 KL 散度计算一个合适的阈值，进行直方图截断，使得截断后的数据较均匀分布在[-127 ~ 127]范围内。

在PTQ后训练量化时，权重参数在训练后就已经确定了，所以权重一般采用非饱和量化方法，也就是最大值量化，激活采用饱和量化方法

权重量化和权重激活量化

根据需要量化的参数可以分类：权重量化和权重激活量化。

• 权重量化：即仅仅需要对网络中的权重执行量化，由于网络的权重一般保存下来，因此可以提前根据权重获得相应的量化参数S和Z，而不需要额外的校准数据集。一般来说，推理过程中，权重值的数量远小于激活值，仅仅对权重执行量化的方法能带来的压缩力度和加速效果都一般
• 权重激活量化：不仅对网络中的权重进行量化，还对激活值进行量化。由于激活值的范围通常不容易提前获得，因此需要在网络推理过程中进行计算或根据模型大致的预测。

量化方法

重点介绍下基于线性量化的对称量化和非对称量化的理论公式

1. 非对称量化
   非对称量化映射实数到b-bit有符号整数 x q ∈ { − 2 b − 1 , − 2 b − 1 + 1 , . . . , 2 b − 1 − 1 } x_q \in \{-2^{b-1},-2^{b-1}+1,...,2^{b-1}-1 \} xq​∈{ −2b−1,−2b−1+1,...,2b−1−1}，非对称函数如下：
   $$\begin{gathered} s=\frac{2^b-1}{\alpha -\beta } \\ z=-round(\beta .s)-2^{b-1} \end{gathered}$$

   其中S是缩放因子，z是零点。在int8量化时，$s=\frac{255}{\alpha -\beta }$和$z=-round(\beta .s)-128$

   为了计算量化的尺度s，需要确定实数的范围，一种很直接的方法就是根据数值范围的大小来确定$\alpha ,\beta$，但这种方法容易受到噪声的影响，比如，有些权重或者激活中可能存在某些离群点，其数值较大，但对结果影响又很小，如果把这些数值也统计，就容易造成浪费。

   我们需要对实数进行截断，同时确保量化后，实数映射到$[-2^{b-1},2^{b-1}-1]$的整数区间

   量化公式如下：
   $$clip(x,l,u)\{\begin{array}{cc}l& x<l\\ x& l\le x\le u\\ u& x>u\end{array}$$
   $$x_q=quantize(x,b,s,z)=clip(round(s.x+z),-2^{b-1},2^{b-1}-1)$$

   反量化操作就是将量化的整数表示转换为实数，公式如下：
   $$\hat{x}=dequantize(x_q,s,z)=\frac{1}{s}(x_q-z)$$

2. 对称量化

   对称量化仅通过缩放变换执行量化范围映射，输入范围和整数范围都是对称的。对于 INT8，使用整数范围 [−127, 127]，而不使用值 -128 以支持对称。TensorRT int8后量化就是选用这种方法。

   实数的可选范围为$[-\alpha ,\alpha ]$，产生一个$b-bit$量化：“
   $$\begin{gathered} s=\frac{2^{b-1}-1}{\alpha } \\ {x}_q=quantize(x,b,s)=clip(round(s.x),-2^{b-1}+1,2^{b-1}-1) \end{gathered}$$

   $\alpha$代表当前输入数据分布中的实数最大值，范围$[-\alpha ,\alpha ]$，$b=8$代表INT8量化。

   对称量化的相应反量化操作如下：

$$\hat{x}=dequantize(x_q,s)=\frac{1}{s}{x}_q$$

量化粒度

在 Tensor 元素之间共享量化参数有多种选择，我们将此选择称为量化粒度。在最粗粒度 per-tensor 上，tensor 中的所有元素共享相同的量化参数。最细的粒度将具有每个元素的单独量化参数。中间粒度在不同维度的张量上重用参数，如对于二维矩阵，每行或每列重用参数；对于三维（类似图像的）张量，每通道重用参数。

在选择粒度时，考虑两个因素：对模型精度和计算成本的影响。为了理解计算成本，我们通过矩阵乘法来说明。考虑一个线性层执行矩阵乘法Y=XW，X为输入W的权重，Y为输出，偏置一般可以去掉，对精度影响不大，所以暂不考虑

输入X的维度为$[m,p]$，W的维度为$[p,n]$，i的范围为$[0,m]，$j的范围为$[0,n]$，k的范围为$[0,p]$，而对应的INT8的输入和权重如下：

• $X=(x_{q,ik})\in {\mathbb{R}}^{m\times p}$为输入激活张量
• $W=(w_{q,kj})\in {\mathbb{R}}^{p\times n}$是权重张量
• $Y=(y_{ij})\in {\mathbb{Z}}^{p\times n}$是输出张量

实数矩阵乘法Y=XW结果可以用量化张量$X_{)}q=(x_{q,ik})\in {\mathbb{R}}^{m\times p}$和$W_q=(w_{q,kj})\in {\mathbb{R}}^{p\times n}$近似

首先对它们进行去量化，然后执行矩阵乘法。考虑采用最细粒度，每个元素量化的对称量化：
$$y_{ij}=\sum _{k=1}^p{x}_{ik}.w_{kj}\approx \sum _{k=1}^{p}dequantize(x_{q,ik},s_{q,ik}).dequantize(w_{q,kj},s_{w,kj})=\sum _{k=1}^p\frac{1}{s_{x,ik}}{x}_{q,ik}.\frac{1}{s_{w,kj}}{w}_{q,kj}$$

为了使用整数矩阵乘法，将缩放因子从上式右边分离出来，因此缩放因子必须与k无关
$$\frac{1}{s_{x,ik}.s_{w,kj}}\sum _{k=1}^p{x}_{q,ik}.w_{q,kj}$$

对于激活量化粒度可以是 per-row 或 per-tensor，对于权重是 per-column 或 per-tensor。
对于激活，出于性能原因，激活应使用 per-tensor 量化粒度。对于线性层，权重应以 per-tenosr 或 per-column 粒度进行量化。卷积中 per-column 对应的粒度是 per-kernel，或等价于per-channel（论文中通常称其为 “per-channel” 权重量化），per-channel 对于每个卷积核具有不同的尺度和偏移。

per-layer 量化：整个张量指定一个量化器（由scale and zero-point 定义）
per-channel ：对于每个卷积核具有不同的尺度和偏移。

实践中权重通常采取per-channel 量化，激活值通常采取per-layer 量化

PTQ和QAT

int8 后量化

在训练后量化中，TensorRT 计算网络中每个张量的比例值。这个过程称为校准，需要您提供有代表性的输入数据，TensorRT 在其上运行网络以收集每个激活张量的统计信息。

根据需要量化的参数可以分类：权重量化和权重激活量化。

• 权重量化：即仅仅需要对网络中的权重执行量化，由于网络的权重一般保存下来，因此可以提前根据权重获得相应的量化参数S和Z，而不需要额外的校准数据集。一般来说，推理过程中，权重值的数量远小于激活值，仅仅对权重执行量化的方法能带来的压缩力度和加速效果都一般
• 权重激活量化：不仅对网络中的权重进行量化，还对激活值进行量化。由于激活值的范围通常不容易提前获得，因此需要在网络推理过程中进行计算或根据模型大致的预测。

如何找到合适的量化参数呢？

对于权重而言，在模型训练完成后数值就基本确定了，而对于特征图来说，却没法事先得知，因此会用一批矫正数据集 (通常就是训练集的一小部分) 跑一遍网络，以此来统计每一层 特征图 的数值范围。

有了权重和特征图的数值范围后，一种很直接方法是根据数值范围的大小确定$r_{max}、r_{min}$，但这种方法很容易受噪声的影响，比如存在某些离群点，其数值较大，但影响小，对结果噪声干扰

需要排除离群点，通常采取一些策略来确定有效的量化参数：

• 指数平滑法：把矫正数据集分为几个batch。逐次输入到网络中统计数值，统计每个batch的量化参数，通过指数平滑法更新更新S和Z参数值

• 直方图截断法：即在计算量化参数Z和S过程中，由于有的特征图会出现偏离较远的奇异值，导致max非常大，可以通过直方图截取的形式，将激活和权重的数值范围统计出一个直方图，根据直方图舍弃前后m%的数据，用剩下的数值来确定min、ma

• 均值和方差：通常会假设 权重和 激活的数值呈正态分布，根据正态分布的性质，可以取区间 ，可以覆盖99%以上的数据。但如果实际分布不是正态分布就可能带来巨大的误差

• KL散度校准法：即通过计算KL散度（也称为相对熵，用以描述两个分布之间的差异）来评估量化前后的两个分布之间存在的差异，搜索并选择KL散度最小的量化参数Z和S作为最终结果。

在介绍Tensort 具体的int8量化原理前，先了解下KL散度的原理

Kullback-Leibler散度

相对熵（relative entropy），又被称为KLD散度（Kullback-Leibler divergence）或信息散度，是两个概率分布（probability distribution）间差异的非对称性度量 。在信息理论中，相对熵等价于两个概率分布的信息熵（Shannon entropy）的差值 。

KL散度的定义：

一个随机变量X的可能值为$X=x_1,x_2,...,x_n$，对应的概率为$p_i=p(X=x_i)$，则随机变量的熵定义为：
$$H(X)=-\sum p(x_i)\log p(x_i)$$

规定当$p(x_I)=0时，p(x_i)\log (p(x_i)=0$)

如果两个随机变量$P,Q$，且其概率分布分别为$p(x),q(x)$，则$p$相对$q$的相对熵为：

$$D_{KL}(p||q)=\sum _{i=1}^{n}P(x)\log \frac{p(x)}{q(x)}$$

针对上述离散变量的概率分布$p(x)、q(x)而言，其交叉熵定义为：

$$H(p,q)=\sum p(x)\log \frac{1}{q(x)}=-\sum p(x)\log q(x)$$
因此，KL散度或相对熵可通过下列式子得出：

$$\begin{gathered} D_{KL}(p||q)=H(p,q)-H(p) \\ =-\sum p(x)\log q(x)-\sum -p(x)\log p(x) \\ =-\sum _{x}p(x)(\log q(x)-\log p(x)) \\ =-\sum _{x}p(x)log\frac{q(x)}{p(x)}=\sum _{x}p(x)log\frac{p(x)}{q(x)} \end{gathered}$$
同样，当P,Q为连续变量时，KL散度的定义为：
$$D(P||Q)=\int P(x)[\log (\frac{P(i)}{Q(i)})]dx$$
在统计学意义上来说，KL散度可以用来衡量两个分布之间的差异程度。若两者差异越小，KL散度越小，反之亦反。当两分布一致时，其KL散度为0。正是因为其可以衡量两个分布之间的差异，所以在VAE、EM、GAN中均有使用到KL散度。

直方图分布

TensoRT使用选用对称量化，量化参数Z=0，只需要计算尺度S，确定KL散度确定截断散度

前面也提到过，直接根据数值范围的大小确定max，易受到噪声的影响，对于一些过大的离群点，其数值较大，影响量化的精度，我们要采用策略进行截断，

不妨先用直方图统计一下的数量具体分布情况。

完成$n_b$个bin的统计直方图histogram，再对histogram进行归一化，即可获得一个概率密度分布$P(i)$
$$P(i)=\frac{histogra{m}_i}{{\sum }_{i=0}^{n_b-1}histogra{m}_i}$$

截断

截断方案：直方图截断，选择从第$T_b$ 个bin截断，$T_b$后面的bin加到$T_{b-1}$ 上，截断后的|x|概率密度分布变为$P_{clip}(i)$

校准流程

TensorRT 采用基于实验的迭代搜索阈值。校准是其中的一个主要部分。校准过程如下所示，具体过程如下：

• 将数据划分为2047个bin
• 取前 128 个bin作为基准，逐次向后搜索，每次扩增一个 bin的长度。
  • 将bin[0],...,bin[i-1]量化为128个bin，计算P的概率分布并归一化
  • 选择从第i个bin截断，i后面的bin加到i-1 上，得到一个基本没有损失的直方图概率分布P_clip，并归一化
  • 计算KL散度
• 选取KL散度最小时对应的calib_amax，calib_amax表示KL散度最小时，直方图对应的``clip边缘位置，该位置clip，效果最好
  

TensoRT采用最简单的办法就是先把数据量化一遍，然后再统计量化后的直方图。但是这样做计算量太大了

INT8线性量化将float32数据量化为int8，那么数据分布也会对应的从个bin被线性压缩到128个bin，所以也可以通过插值的方式直接缩小P计算出分布Q，更多参考：https://zhuanlan.zhihu.com/p/387072703

代码实现可以参考pytorch_quantization库中量化尺度的entorpy校准方案：
代码位于pytorch_quantization/tools/pytorch-quantization/pytorch_quantization/calib/calibrator.py

_compute_amax_entropy：即通过计算KL散度（也称为相对熵，用以描述两个分布之间的差异）来评估量化前后的两个分布之间存在的差异，最小化量化后int8与原始float32数据之间的信息损失，遍历搜索并选择KL散度最小的calib_amax

def _compute_amax_entropy(calib_hist, calib_bin_edges, num_bits, unsigned, stride=1, start_bin=128):
    # 返回使收集的直方图的KL散度最小化的amax
    """Returns amax that minimizes KL-Divergence of the collected histogram"""

    # If calibrator hasn't collected any data, return none
    if calib_bin_edges is None and calib_hist is None:
        return None

    def _normalize_distr(distr):
        summ = np.sum(distr)
        if summ != 0:
            distr = distr / summ

    bins = calib_hist[:]
    bins[0] = bins[1]

    total_data = np.sum(bins)

    divergences = []
    arguments = []

    # we are quantizing to 128 values + sign if num_bits=8
    nbins = 1 << (num_bits - 1 + int(unsigned)) # 对称量化 nbins=128

    starting = start_bin
    stop = len(bins) # 4028

    new_density_counts = np.zeros(nbins, dtype=np.float64)

    #首次遍历 i=128
    for i in range(starting, stop + 1, stride):
        new_density_counts.fill(0) 
        # 这里是先进行量化，再计算数据分布Q,耗时比较大
        # 把bin[0],...,bin[i-1]量化为128个bin
        space = np.linspace(0, i, num=nbins + 1)
        # numpy.digitize(array_x, bins, right=False)：返回array中每一个值在bins中所属的位置
        # 记录量化前的i在量化后的位置
        digitized_space = np.digitize(range(i), space) - 1
        digitized_space[bins[:i] == 0] = -1 # 直方图值为0 对应 digitized_space 值取-1

        # 计算量化后的数据分布Q
        for idx, digitized in enumerate(digitized_space):
            if digitized != -1:
                # 将直方图柱子不为0依次累加
                new_density_counts[digitized] += bins[idx]

        counter = Counter(digitized_space) # Counter：统计可迭代序列中每个元素出现次数
        for key, val in counter.items():
            if key != -1:
                # 计算分布Q:new_density_counts
                new_density_counts[key] = new_density_counts[key] / val

        new_density = np.zeros(i, dtype=np.float64)
        # 剔除直方图值为0的
        for idx, digitized in enumerate(digitized_space):
            if digitized != -1:
                new_density[idx] = new_density_counts[digitized]

        total_counts_new = np.sum(new_density) + np.sum(bins[i:])
        # 归一化
        _normalize_distr(new_density)

        # 取前i个bin
        reference_density = np.array(bins[:len(digitized_space)])
        # 选择从第i个bin截断，i后面的bin加到i-1 上，得到一个基本没有损失的直方图P_clip
        # reference_density 代表原始float数据截断后的分布情况
        reference_density[-1] += np.sum(bins[i:])

        total_counts_old = np.sum(reference_density)
        if round(total_counts_new) != total_data or round(total_counts_old) != total_data:
            raise RuntimeError("Count mismatch! total_counts_new={}, total_counts_old={}, total_data={}".format(
                total_counts_new, total_counts_old, total_data))

        _normalize_distr(reference_density)
        # 计算KL散度，散度越小代表分布越相似
        ent = entropy(reference_density, new_density)
        divergences.append(ent)
        arguments.append(i)

    divergences = np.array(divergences)
    logging.debug("divergences={}".format(divergences))
    last_argmin = len(divergences) - 1 - np.argmin(divergences[::-1])
    # calib_amax代表截断float的最大值，在此处截断，量化效果最好
    calib_amax = calib_bin_edges[last_argmin * stride + starting]
    calib_amax = torch.tensor(calib_amax.item()) 

    return calib_amax

感知训练量化

如果PTQ中模型训练和量化是分开的，而QAT则是在模型训练时加入了伪量化节点，用于模拟模型量化时引起的误差。以INT8量化为例，QAT处理流程为

• 首先在数据集上以FP32精度进行模型训练，得到训练好的baseline模型；
• 在baseline模型中插入伪量化节点，得到QAT模型，并且在数据集上对QAT模型进行finetune；
• 伪量化节点会模拟推理时的量化过程并且保存finetune过程中计算得到的量化参数；
• finetune完成后，使用3. 中得到的量化参数对QAT模型进行量化得到INT8模型，并部署至推理框架中进行推理

QAT方式需要重新对插入节点之后的模型进行finetune，通过伪量化操作，可以是网络各层的weights和activation输出分布更加均匀，相对于PTQ可以获得更高的精度。

伪量化节点(Fake-quantize)

那么什么是伪量化呢？伪量化实际上是quantization+dequantization的结合，实际上就是模拟量化round引起的误差

$$clamp(r,a,b)=min(Max(x,a),b)$$

$$s(a,b,n)=\frac{b-a}{n-1}$$

$$q(r,a,b,n)=[\frac{clamp(r,a,b)-a}{s(a,b,n)}]s(a,b,n)+a$$

伪量化的操作看起来输入输出没变，但是实际上在其中模拟了量化round操作，将这种误差当做一种训练的噪声，在QAT finetune的同时，模型会去适应这种噪声，从而在最后量化为INT8时，减少精度的损失。

伪量化节点插入位置就是需要进行量化操作的位置，论文中在weights输入conv之前（weight quantization）以及activation之后(activation quantizaion)插入了伪量化节点，QAT训练时所有计算都使用FP32单精度。

QAT finetune过程量化参数处理

前面提到伪量化节点会保存finetune过程中的量化参数，伪量化节点的计算公式中$[a,b]$ 即为FP32浮点数值的范围，这个值将会在finetune过程中进行估计与更新，上面介绍了伪量化节点分别weight quantization以及activation quantizaion：

• 对于weight quantization的量化节点，直接将$[a,b]$设置为weights的最大值与最小值；
• 对于activation quantizaion，处理方式类似于batch norm，使用了指数移动平均，在finetune的每个batch动态地更新$[a,b]$ ；
  最后量化模型的时候，只需设置$S=s(a,b,n)$、$Z=z(a,b,n)$ 即可。

参考：

• https://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf

• https://www.nvidia.com/en-us/on-demand/session/gtcspring21-s31653/

• https://zhuanlan.zhihu.com/p/387072703

• Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference

• https://www.nvidia.com/en-us/on-demand/session/gtcsj20-s21664/