TensorFlow Lite 采坑记（二）：Model quantization 模型量化

前言

从2012年AlexNet夺得当年ImageNet冠军开始，深度学习就开始呈爆破式发展，这个过程中迭代出很多优秀的深度学习代码框架，随着技术逐渐成熟，工业界对技术落地的需求越来越迫切，本文以TensorFlow为例，探究模型量化的相关原理和方法。

为什么要做模型量化

深度学习已被证明在包括图像分类（Image Classification），目标检测（Object Detection），自然语言处理（Natural Language Processing）等任务上效果很好。但到了实践中，模型的复杂程度度往往限制了其在各种现实应用场景下的部署。
图一：模型 MAC和网络预测准确度。

自 AlexNet 伊始，基于 ImageNet 的深度学习算法（或模型）改进都和模型大小相关。在如上的由 Google 研究人员给出的 图一 中，垂直方向是网络在给定任务上的效果，横向是网络的大小。通过对比图中 AlexNet，GoogleNet 和 VGG 的趋势不难看出，当模型变大时网络的准确度可能更高。

随着模型预测越来越准确，网络越来越深，神经网络消耗的内存大小成为问题（见下图），尤其是在移动设备上。通常情况下，目前（2019年初）的手机一般配备 4GB 内存来支持多个应用程序的同时运行。而三个模型运行一次通常就要占用1GB内存。

模型大小不仅是内存容量问题，也是内存带宽问题。模型在每次预测时都会使用模型的权重，图像相关的应用程序通常需要实时处理数据，这意味着至少 30 FPS。因此，如果部署相对较大的 ResNet-50 网络来分类，运行网络模型就需要 3GB/s 的内存带宽。网络运行时，内存、CPU 和电池会都在飞速消耗。

1. 量化原理

总的来说， 量化桥接了浮点和定点，使模型大幅度压缩同时使精度损失保持在合理范围内，减小了带宽需求。值得一提的是，量化后的模型在做inference的时候还是会将权重转换为浮点数。原理参考文献：Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference
以下是基于该文献对量化的一些理解。

1.1 Quantization Scheme
定义$r$为浮点类型的实际值， 定义$q$为整型的量化值，量化模式可以简述为：
$$q=\frac{r}{S}+Z$$
反之：
$$r=S(q-Z)$$
其中$$S=\frac{r{}_{max}-r{}_{min}}{q{}_{max}-q{}_{min}}$$
$$Z=q{}_{max}-r{}_{max}÷S$$
这里的 $S$ 和 $Z$ 均为量化参数，前者如字面意思所示，$Z$ 表示浮点数的 0 量化后对应的整型值，由于 0 在神经网络中有着特殊的含义（padding等），故必须有精确的整型值对应 0。对于量化后的值 $q$，通过量化参数 $(S,Z)$ 可恢复到所代表的浮点数值。
假设使用8bit量化，量化后的整型值应该在[0，255]之间，对应的可以表示的浮点数值范围：
$$min=S(0-Z)$$
$$max=S(255-Z)$$
对于超出[min, max]范围的浮点数值，量化后必须做截断。对于网络的每一层（或者每个input array），都会有不同的min、max，因为是由不同的量化参数$S$， $Z$决定的。在[min, max]范围的浮点数经过量化后可以得到uint8的8bit整数。

1.2 Integer-arithmetic-only matrix multiplication
解决了数值的量化，接下来就是计算层面的问题。因为直观上来看量化是将权重和神经网络中的数据流从float转换成了uint，但在同时需要控制在计算时的精度损失在一个合理范围内，否则量化将毫无意义。可以预见的是，在整个神经网络的计算过程中矩阵乘法占了相当大的比重。
考虑两个$N$x$N$矩阵$r_1$和$r_2$的乘法，假设$r_3$ = $r_1$$r_2$，进行INT8量化。 易知：
$$r{}_a^{(i,j)}=S{}_a(q{}_a^{(i,j)}-Z{}_a)$$

矩阵乘法可改写为：

移项：

where $$M=\frac{S{}_{1}S{}_2}{S{}_3}=\frac{M{}_0}{2{}^n}$$
易知式中$M$为FP32，其他均为INT8。实际上两个INT8之间的算术会累加到INT16或INT32，即超出INT8值域，但$Z{}_{1-3}$项会使大多数的计算结果落入INT8值域中。
除乘法外，还有量化加法、减法、除法和指数等等运算。采用这些方法，量化神经网络可以在保证精度的同时大大减小模型大小和所需带宽。

2. 模型量化加速原理

上节探究了量化使模型大小大幅度下降并保持了一定的准确性，其实理论上来说，量化后的模型在CPU上的计算速度也会得到提升。

2.1 CPU命中率与L1 Cache
http://cpu.zol.com.cn/31/311395_all.html
2.2 im2col 与 Gemm
https://jackwish.net/gemm-optimization-and-convolution.html
https://petewarden.com/2015/04/20/why-gemm-is-at-the-heart-of-deep-learning/

相信读完这几篇文章就会对其有一定理解了。

3. 两种方法

3.1 Post-training quantization
官方文档提供了其python API和相关例子，详见：https://www.tensorflow.org/lite/performance/post_training_quantization

3.2 Quantization-aware training
TensorFlow 在支持训练后量化的同时还引入了训练时量化 Quantization-aware Training https://github.com/tensorflow/tensorflow/tree/r1.13/tensorflow/contrib/quantize ，其中包括四个步骤：

1. 用常规方法训练一个 TensorFlow 浮点模型。
2. 用 tf.contrib.quantize 重写网络以插入Fake-Quant 节点并训练 min/max。
3. 用 TensorFlow Lite 工具量化网络（该工具读取步骤 2 训练的 min/max）
4. 用 TensorFlow Lite 部署量化的网络。

步骤 2 即所谓的量化感知训练（Quantization-aware Training），其中网络的前向（forward）模拟 INT8 计算，反向（backward）仍然是 FP32 。下图左半部分是量化网络，它接收 INT8 输入和权重并生成 INT8 输出。图十二右半部分是步骤 2 重写的网络，其中 Fake-Quant 节点（粉色）在训练期间将 FP32 张量量化为 INT8 （严格来讲仍然是经过 Quantize/Dequantize 过程的 FP32）

细节请参考论文Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference

3.3 流程图

流程图是基于TensorFlow Lite官方文档整理，笔者分别尝试了Convert a TensorFlow GraphDef for quantized inference（命令行）和 post-training quantization两种方法。
官方文档链接：
Convert a TensorFlow GraphDef for quantized inference:
https://www.tensorflow.org/lite/convert/cmdline_examples#convert_a_tensorflow_graphdef_for_quantized_inference_
“dummy-quantization”：
https://www.tensorflow.org/lite/convert/cmdline_examples#use_dummy-quantization_to_try_out_quantized_inference_on_a_float_graph_
post-training quantization：
https://www.tensorflow.org/lite/performance/post_training_quantization
quantization-aware training：
https://github.com/tensorflow/tensorflow/tree/r1.13/tensorflow/contrib/quantize

4. limitations

主要有两点：

1. 量化后的模型暂不支持GPU或NNAPI加速。
2. 对于不同的CPU架构，其表现可能有较大差别（待验证）。

5. Quantization model 和 Float model 比较（数据请求中）

Reference：
https://jackwish.net/neural-network-quantization-introduction-chn.html
https://sf-zhou.github.io/ml/tensorflow_lite_and_uint8_quantization.html
https://zhuanlan.zhihu.com/p/42811261