基于OneFlow实现量化感知训练

Posted just_sort

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了基于OneFlow实现量化感知训练相关的知识，希望对你有一定的参考价值。

【GiantPandaCV导语】本文介绍了量化感知训练的原理，并基于OneFlow实现了一个量化感知训练Demo，并介绍了在具体实现中的各种细节。希望对想学习量化感知训练的读者有用，本文仅做学习交流。

0x0. 前言

这篇文章主要是讲解一下量化感知训练的原理，以及基于OneFlow实现一个Demo级别的手动量化感知训练。

0x1. 后量化以及量化感知训练原理

这里说的量化一般都是指的Google TFLite的量化方案，对应的是Google 的论文 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference。虽然TfLite这套量化方案并不是很难，但在实际处理的时候细节还是比较多，一时是很难说清楚的。

所以，这里推荐一系列讲解TFLite后量化和量化感知训练原理的文章，看一下这几篇文章阅读本文就没有任何问题了。

这里我简单的总结一下，无论是TFLite的量化方案还是TensorRT的后量化方案，他们都会基于原始数据和量化数据的数值范围算出一个缩放系数scale和零点zero_point，这个zero_point有可能是0（对应对称量化），也有可能不是0（对应非对称量化）。然后原始数据缩放之后减掉零点就获得了量化后的数据。这里的关键就在于缩放系数scale和zero_point怎么求，Google的TFLite使用下面的公式：

S = r m a x − r m i n q m a x − q m i n S = \\frac{r_{max}-r_{min}}{q_{max}-q_{min}} S=qmax​−qmin​rmax​−rmin​​

Z = r o u n d ( q m a x − r m a x S ) Z = round(q_{max} - \\frac{r_{max}}{S}) Z=round(qmax​−Srmax​​)

其中， r r r表示浮点实数， q q q表示量化后的定点整数， r m a x r_{max} rmax​和 r m i n r_{min} rmin​分别是 r r r的最大值和最小值， q m i n q_{min} qmin​和 q m a x q_{max} qmax​表示 q q q的最大值和最小值，如果是有符号8比特量化，那么 q m i n = − 128 q_{min}=-128 qmin​=−128， q m a x = 127 q_{max}=127 qmax​=127，如果是无符号那么 q m i n = 0 q_{min}=0 qmin​=0， q m a x = 255 q_{max}=255 qmax​=255。 S S S就代表scale， Z Z Z就代表zero_point。

要求取scale和zero_point关键就是要精确的估计原始浮点实数的最大值和最小值，有了原始浮点实数的最大值和最小值就可以代入上面的公式求出scale和zero_point了。所以后训练量化以及量化感知训练的目的是要记录各个激活特征图和权重参数的scale和zero_point。

在后训练量化中，做法一般是使用一部分验证集来对网络做推理，在推理的过程中记录激活特征图以及权重参数的最大和最小值，进而计算scale和zero_point。而量化感知训练则在训练的过程中记录激活特征图和权重参数的最大和最小值来求取scale和zero_point。量化感知训练和后训练量化的主要区别在于它会对激活以及权重做模拟量化操作，即FP32->INT8->FP32。这样做的好处是可以模拟量化的实际运行过程，将量化过程中产生的误差也作为一个特征提供给网络学习，一般来说量化感知训练会获得比后训练量化更好的精度。

0x2. 组件

在上一节中主要提到了记录激活和权重的scale和zero_point，以及模拟量化，量化这些操作。这对应着三个量化训练中用到的三个基本组件，即MinMaxObserver，FakeQuantization，Quantization。下面我们分别看一下在OneFlow中这三个组件的实现。

组件1. MinMaxObserver

从这个文档我们可以看到MinMaxObserver操作被封装成oneflow.nn.MinMaxObserver这个Module（Module在Pytorch中对应torch.nn.Module，然后OneFlow的接口也在靠近Pytorch，也对应有oneflow.nn.Module，因此这里将其封装为oneflow.nn.Module）。这个Module的参数有：

  • quantization_bit表示量化Bit数
  • quantization_scheme 表示量化的方式，有对称量化symmetric和非对称量化affine两种，区别就是对称量化浮点0和量化空间中的0一致
  • quantization_formula 表示量化的方案，有Google和Cambricon两种，Cambricon是中科寒武纪的意思
  • per_layer_quantization 表示对当前的输入Tensor是PerChannel还是PerLayer量化，如果是PerLayer量化设置为True。一般激活特征图的量化都是PerLayer，而权重的量化可以选择PerLayer或者PerChannel。

下面看一下在Python层的用法：

>>> import numpy as np
>>> import oneflow as flow

>>> weight = (np.random.random((2, 3, 4, 5)) - 0.5).astype(np.float32)

>>> input_tensor = flow.Tensor(
...    weight, dtype=flow.float32
... )

>>> quantization_bit = 8
>>> quantization_scheme = "symmetric"
>>> quantization_formula = "google"
>>> per_layer_quantization = True

>>> min_max_observer = flow.nn.MinMaxObserver(quantization_formula=quantization_formula, quantization_bit=quantization_bit,
... quantization_scheme=quantization_scheme, per_layer_quantization=per_layer_quantization)

>>> scale, zero_point = min_max_observer(
...    input_tensor, )

在设定好相关量化配置参数后，传入给定Tensor即可统计和计算出该设置下的Tensor的scale和zero_point。

上面讲的是Python前端的接口和用法，下面看一下在OneFlow中这个Module的具体实现，我们以CPU版本为例（GPU和CPU的Kernel实现是一致的），文件在oneflow/user/kernels/min_max_observer_kernel.cpp，核心实现是如下三个函数：

// TFLite量化方案，对称量化
template<typename T>
void GenQuantScaleSymmetric(const T* in_ptr, const int32_t quantization_bit,
                            const int64_t num_elements, T* scale, T* zero_point) {
  T in_max = *std::max_element(in_ptr, in_ptr + num_elements);
  T in_min = *std::min_element(in_ptr, in_ptr + num_elements);

  in_max = std::max(std::abs(in_max), std::abs(in_min));

  T denominator = static_cast<T>(pow(2.0, quantization_bit - 1)) - 1;

  *scale = in_max / denominator;
  *zero_point = 0;
}

// TFLite量化方案，非对称量化
template<typename T>
void GenQuantScaleAffine(const T* in_ptr, const int32_t quantization_bit,
                         const int64_t num_elements, T* scale, T* zero_point) {
  T in_max = *std::max_element(in_ptr, in_ptr + num_elements);
  T in_min = *std::min_element(in_ptr, in_ptr + num_elements);

  T denominator = static_cast<T>(pow(2.0, quantization_bit)) - 1;

  *scale = (in_max - in_min) / denominator;
  *zero_point = -std::nearbyint(in_min / (*scale));
}

//寒武纪量化方案
template<typename T>
void GenQuantScaleCambricon(const T* in_ptr, const int32_t quantization_bit,
                            const int64_t num_elements, T* scale, T* zero_point) {
  T in_max = *std::max_element(in_ptr, in_ptr + num_elements);
  T in_min = *std::min_element(in_ptr, in_ptr + num_elements);

  in_max = std::max(std::abs(in_max), std::abs(in_min));

  *scale = std::floor(std::log2(in_max)) - (quantization_bit - 2);
  *zero_point = 0;
}

除了这三个函数之外，另外一个关键点就是对per_layer_quantization参数的处理了，逻辑如下：

如果是PerChannel量化则对每个输出通道求一个scale和zero_point。想了解更多PerLayer量化以及PerChannel量化的知识可以看这篇文章：神经网络量化–per-channel量化 。

组件2：FakeQuantization

同样，FakeQuantization也被封装为一个oneflow.nn.Module。在上一节提到，量化感知训练和后训练量化的主要区别在于它会对激活以及权重参数做模拟量化操作，即FP32->INT8->FP32。通过这种模拟将量化过程中产生的误差也作为一个特征提供给网络学习，以期在实际量化部署时获得更好的准确率。这个接口有以下参数：

  • scale：由MinMaxObserver组件算出来的量化scale
  • zero_point：由MinMaxObserver组件算出来的量化zero_point
  • quantization_bit： 量化比特数
  • quantization_scheme 表示量化的方式，有对称量化symmetric和非对称量化affine两种，区别就是对称量化浮点0和量化空间中的0一致
  • quantization_formula 表示量化的方案，有Google和Cambricon两种，Cambricon是中科寒武纪的意思

Python层的示例用法如下：

>>> import numpy as np
>>> import oneflow as flow

>>> weight = (np.random.random((2, 3, 4, 5)) - 0.5).astype(np.float32)

>>> input_tensor = flow.Tensor(
...    weight, dtype=flow.float32
... )

>>> quantization_bit = 8
>>> quantization_scheme = "symmetric"
>>> quantization_formula = "google"
>>> per_layer_quantization = True

>>> min_max_observer = flow.nn.MinMaxObserver(quantization_formula=quantization_formula, quantization_bit=quantization_bit,
... quantization_scheme=quantization_scheme, per_layer_quantization=per_layer_quantization)
>>> fake_quantization = flow.nn.FakeQuantization(quantization_formula=quantization_formula, quantization_bit=quantization_bit,
... quantization_scheme=quantization_scheme)

>>> scale, zero_point = min_max_observer(
...    input_tensor,
... )

>>> output_tensor = fake_quantization(
...    input_tensor,
...    scale,
...    zero_point,
... )

在执行FakeQuantizaton必须知道输入Tensor的scale和zero_point，这是由上面的MinMaxObserver组件获得的。

接下来看一下FakeQuantization组件C++层的实现，仍然有三个核心函数：

// TFLite量化方案，对称量化
template<typename T>
void FakeQuantizationPerLayerSymmetric(const T* in_ptr, const T scale,
                                       const int32_t quantization_bit, const int64_t num_elements,
                                       T* out_ptr) {
  T upper_bound = static_cast<T>(pow(2.0, quantization_bit - 1)) - 1;
  T lower_bound = -upper_bound - 1;
  FOR_RANGE(int64_t, i, 0, num_elements) {
    T out = std::nearbyint(in_ptr[i] / scale);
    out = out > upper_bound ? upper_bound : out;
    out = out < lower_bound ? lower_bound : out;
    out_ptr[i] = out * scale;
  }
}

// TFLite量化方案，非对称量化
template<typename T>
void FakeQuantizationPerLayerAffine(const T* in_ptr, const T scale, const T zero_point,
                                    const int32_t quantization_bit, const int64_t num_elements,
                                    T* out_ptr) {
  T upper_bound = static_cast<T>(pow(2.0, quantization_bit)) - 1;
  T lower_bound = 0;
  uint8_t zero_point_uint8 = static_cast<uint8_t>(std::round(zero_point));
  FOR_RANGE(int64_t, i, 0, num_elements) {
    T out = std::nearbyint(in_ptr[i] / scale + zero_point_uint8);
    out = out > upper_bound ? upper_bound : out;
    out = out < lower_bound ? lower_bound : out;
    out_ptr[i] = (out - zero_point_uint8) * scale;
  }
}
// 寒武纪量化方案
template<typename T>
void FakeQuantizationPerLayerCambricon(const T* in_ptr, const T shift,
                                       const int32_t quantization_bit, const int64_t num_elements,
                                       T* out_ptr) {
  T upper_bound = static_cast<T>(pow(2.0, quantization_bit - 1)) - 1;
  T lower_bound = -upper_bound - 1;
  T scale = static_cast<T>(pow(2.0, static_cast<int32_t>(shift)));
  FOR_RANGE(int64_t, i, 0, num_elements) {
    T out = std::nearbyint(in_ptr[i] / scale);
    out = out > upper_bound ? upper_bound : out;
    out = out < lower_bound ? lower_bound : out;
    out_ptr[i] = out * scale;
  }
}

需要注意的一点是由于FakeQuantization要参与训练，所以我们要考虑梯度怎么计算？从上面的三个核心函数实现中我们可以发现里面都用了std::nearbyint函数，这个函数其实就对应numpy的round操作。而我们知道round函数中几乎每一处梯度都是0，所以如果网络中存在这个函数，反向传播的梯度也会变成0。

因此为了解决这个问题，引入了Straight Through Estimator。即直接把卷积层（这里以卷积层为例子，还包含全连接层等需要量化训练的层）的梯度回传到伪量化之前的weight上。这样一来，由于卷积中用的weight是经过伪量化操作的，因此可以模拟量化误差，把这些误差的梯度回传到原来的 weight，又可以更新权重，使其适应量化产生的误差，量化训练也可以正常运行。

具体的实现就非常简单了，直接将dy赋值给dx，在OneFlow中通过identity这个Op即可：

组件三：Quantization

上面的FakeQuantization实现了FP32->INT8->FP32的过程，这里还实现了一个Quantization组件备用。它和FakeQuantization的区别在于它没有INT8->FP32这个过程，直接输出定点的结果。所以这个组件的接口和C++代码实现和FakeQuantization基本完全一样（反向就不需要了），这里不再赘述。之所以要独立这个组件是为了在训练完模型之后可以将神经网络的权重直接以定点的方式存储下来。后面的Demo中将体现这一点。

0x3. 基于OneFlow量化感知训练AlexNet

下面以AlexNet为例，基于OneFlow的三个量化组件完成一个量化感知训练Demo。这里先贴一下实验结果：

训练的数据集是ImageNet的一个子集，详细信息可以https://github.com/Oneflow-Inc/models/pull/78看到。在8Bit的时候无论是选用Google还是寒武纪，对称还是非对称，PerLayer还是PerChannel，量化感知训练后的模型精度没有明显降低。一旦将量化Bit数从8降到4，在相同的超参配置下精度有了明显下降。

下面分享一下这个基于OneFlow的量化感知训练Demo的做法：

首先代码结构如下：

- quantization
	- quantization_ops 伪量化OP实现
	    - q_module.py 实现了Qparam类来管理伪量化参数和操作和QModule基类管理伪量化OP的实现
	    - conv.py 继承QModule基类，实现卷积的伪量化实现
	    - linear.py 继承QModule基类，实现全连接层的伪量化实现
	    - ...
	- models 量化模型实现
	    - q_alexnet.py 量化版AlexNet模型
	- quantization_aware_training.py 量化训练实现
	- quantization_infer.py 量化预测实现
	- train.sh 量化训练脚本
	- infer.sh 量化预测脚本

  • 由于量化训练时需要先统计样本以及中间层的 scale、zeropoint，同时也频繁涉及到一些量化、反量化操作，所以实现一个QParam基类封装这些功能。

  • 实现了一个量化基类QModule，提供了三个成员函数__init__，freeze。

    • __init__函数除了需要i指定quantization_bit，quantization_scheme，quantization_formula，per_layer_quantization参数外，还需要指定是否提供量化输入参数(qi) 及输出参数 (qo)。这是因为不是每一个网络模块都需要统计输入的 scale，zero_point，大部分中间层都是用上一层的qo来作为自己的qi，另外有些中间层的激活函数也是直接用上一层的 qi来作为自己的qi和qo。
    • freeze 这个函数会在统计完 scale，zero_point 后发挥作用，这个函数和后训练量化和模型转换有关。如下面的量化公式所示，其中很多项是可以提前计算好的，freeze 就是把这些项提前固定下来，同时也将网络的权重由浮点实数转化为定点整数。

  • 基于这个QModule基类定义QConv2d，QReLU，QConvBN等等。

QConvBN表示Conv和BN融合后再模拟量化。原理可以看第一节的第4篇参考资料。这里以QConv2d为例看看它的实现：

import oneflow as flow
from quantization_ops.q_module import QModule, QParam

__all__ = ["QConv2d"]


class QConv2d(QModule):

    def __init__(self, conv_module, qi=True, qo=True, quantization_bit=8, quantization_scheme='symmetric', quantization_formula='google', per_layer_quantization=True):
        super(QConv2d, self).__init__(qi=qi, qo=qo, quantization_bit=quantization_bit, quantization_scheme=quantization_scheme,
                                      quantization_formula=quantization_formula, per_layer_quantization=per_layer_quantization)
        self.quantization_bit = quantization_bit
        self.quantization_scheme = quantization_scheme
        self.quantization_formula = quantization_formula
        self.per_layer_quantization = per_layer_quantization
        self.conv_module = conv_module
        self.fake_quantization = flow.nn.FakeQuantization(
            quantization_formula=quantization_formula, quantization_bit=quantization_bit, quantization_scheme=quantization_scheme)
        self.qw = QParam(quantization_bit=quantization_bit, quantization_scheme=quantization_scheme,
                         quantization_formula以上是关于基于OneFlow实现量化感知训练的主要内容，如果未能解决你的问题，请参考以下文章
 

 深度学习框架量化感知训练的思考及OneFlow的一种解决方案


 适配PyTorch FX，OneFlow让量化感知训练更简单


 深度学习框架量化感知训练的思考及OneFlow的解决方案


 基于pytorch量化感知训练（mnist分类）--浮点训练vs多bit后量化vs多bit量化感知训练效果对比


 开源项目｜基于darknet实现量化感知训练，已实现yolov3-tiny所有算子


 基于pytorch后量化（mnist分类）---浮点训练vs多bit后量化vs多bit量化感知训练效果对比