TensorFlow模型量化

封面图片：Photo by Daniel von Appen on Unsplash

HRNet人脸特征点检测模型的实现是在PC下完成的。但是实际生活中在手机上使用该功能的场景数不胜数：人脸解锁、人脸支付、短视频与直播人脸特效都能看到它的身影。与PC不同，手持设备在运算能力、存储容量、电能供应等方面无法与数百瓦特功率的台式机相比。因此，深度学习模型在移动终端运行需要经过特别的改造。

TensorFlow当前提供了4种模型优化策略，他们是：

• 训练后量化
• 训练感知量化
• 剪枝
• 权重聚类

本文内容涉及第一种方式——训练后量化。并以HRNet模型为例，展示了量化过程以及结果。

TensorFlow文档提供了有关量化的详细内容，推荐阅读。

Post-training quantization | TensorFlow Lite

TensorFlow

训练后量化的五种方式

量化的本质是降低模型权重参数与激活函数的精度，进而降低了运算量与存储负担。例如HRNet在构建与训练时模型的参数精度为32位浮点数（FP32）。训练后量化可以将已经训练好的模型数据类型替换为16位浮点数（FP16）甚至8位整数（INT8）。按照具体的量化规格与对象的不同，还可以再细分为5类。

动态范围量化

该方式将模型权重精度转换为INT8。“动态范围”算子会根据激活函数的范围动态的将其转换为8bit整数，实现与8bit权重的运算。不过运算最终的输出还是以浮点数的形式保存。该方式的加速效果不及定点运算。实现代码如下：

# 构建模型转换器
converter = tf.lite.TFLiteConverter.from_saved_model("./exported")

# 设定量化策略
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 执行转换
tflite_model = converter.convert()

量化前saved_model模型体积为19.8MB，量化后为3.4MB，缩小了82%。

全整数量化

模型全部以INT8整型的方式运算。但是按照输入输出的数据精度不同，细分两种：输入输出为浮点数；输入输出也是整数。

浮点输入输出

该方式需提供示例数据以便定量输入与激活函数的动态范围。

# 定义示例数据生成器
def representative_dataset():
    for sample in samples:
        yield [sample.image]

# 为转换器提供示例数据
converter.representative_dataset = representative_dataset

量化后的模型体积为3.6MB，缩小了81%。

INT8输入输出

与上边一种方式的根本差异在于输入与输出也是INT8。某些特殊的硬件设备只支持INT8运算，如Google的Edge TPU。

# 设定支持的算子类型为INT8
converter.target_spec.supported_ops = [
            tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

# 设定输入输出类型为INT8
converter.inference_input_type = tf.int8 
converter.inference_output_type = tf.int8 

量化后的模型体积为3.4MB，缩小了82%。

16位浮点数量化

顾名思义，该方式将权重与激活函数均转换为16位浮点数。相比32位浮点数，精度下降的并不多，这意味着模型准确率不会降低太多。从硬件角度分析部分GPU可以提供16位浮点数运算，这可以加速模型推演。

# 指定支持数据类型
converter.target_spec.supported_types = [tf.float16]

量化后的模型体积为6.6MB，比32位模型缩小了66%。

16×8量化

激活函数量化为16bit浮点数，权重量化为8bit整型，而bias量化为64bit整型。 这种方式可以提升超分辨率、图像降噪、HDR重建的性能。注意TensorFlow当前提供的转换仍处于实验阶段。

converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8]

不过该量化方式支持的算子有限。HRNet的转换失败。

准确率比较

量化的副作用在于准确率下降。以下是HRNet采用不同方式转换后的准确率分析。

所有转换后的模型NME都存在不同程度的上升。

表格中没有浮点输入输出的全整数量化数据。因为该模型在PC上推演时速度急剧下降，2500个样本可能需要好几个小时才能完成。具体可以参考下github上的这个issue。

Very slow quantized tflite model · Issue #40183 · tensorflow/tensorflow

System information OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04 TensorFlow installed from (source or binary): binary TensorFlow version (or github SHA if from source)...

GitHubtensorflow

公众号发送“量化”，可以获得适用于HRNet的完整量化转换代码。