TensorFlow 模型优化工具包 — 训练时量化

2020 年 4 月 15 日 TensorFlow

文 / TensorFlow 模型优化团队

我们很高兴发布训练时量化 (QAT) API，这也是 TensorFlow 模型优化工具包中的一部分。借助 QAT，您能够利用量化在性能和大小上的优势训练和部署模型，同时保持与最初接近的准确率。这项工作是我们发展蓝图中的一部分，旨在支持开发更小、更快的机器学习 (ML) 模型。如需了解更多知识背景，请参阅过去有关训练后量化、float16 量化和稀疏性（剪枝）的文章。

TensorFlow 模型优化工具包
http://tensorflow.google.cn/model_optimization
发展蓝图
https://tensorflow.google.cn/model_optimization/guide/roadmap

量化会导致损失

量化是以一定精度损失为代价，减少 ML 模型的参数与运算量的方法。此过程可提高模型的执行性能和效率。例如，TensorFlow Lite 8 位整数量化产生的模型大小最多可缩小至四分之一，计算速度可提高至 1.5~4 倍，且 CPU 功耗更低。此外，量化可实现在专用的神经网络加速器（如 Coral 中的 Edge TPU）上运行模型，但支持的数据类型通常有限。

TensorFlow Lite
https://tensorflow.google.cn/lite/performance/quantization_spec

但是，从高精度转换为低精度的过程其实会导致损失。如下图所示，量化将一小部分浮点值压缩到固定数量的信息存储分区 (information bucket) 中。

小范围的 float32 值映射到 int8 是有损转换，因为 int8 只有 255 个信息通道

此过程会导致信息损失。模型的参数（或权重）现只能取用一小部分值，同时值之间的微小差异会丢失。例如，[2.0, 2.3] 范围内的所有值现在都只能在一个存储分区中表示。类似于在用整数表示小数时的舍入误差。

还有其他造成损失的原因。如在多个乘加计算中使用这些有损值时，损失将会叠加。而且，在下一次计算时，需要将累加为 int32 整数的 int8 值重新转换为 int8 值，从而会引入更多计算误差。

训练时量化

QAT 的核心思想是，可以在训练过程的正向传递中模拟低精度的推理时间计算。这项研究成果要归功于 Google 移动视觉团队成员 Skirmantas Kligys 的原始创新。它将量化误差作为训练过程中的干扰并计入整体损失，而优化算法会尽量将它降至最低。因此，该模型可学习对量化更稳健的参数。

研究成果
https://arxiv.org/abs/1712.05877

如果训练不适用，请查看训练后量化，这也是 TensorFlow Lite 模型转换的一部分。如果研究人员和硬件设计人员想要尝试多种量化策略（超出 TensorFlow Lite 支持范围的策略）和/或模拟量化对不同硬件后端准确率的影响，那么 QAT 对其也非常有用。

训练后量化
https://tensorflow.google.cn/lite/performance/post_training_quantization
TensorFlow Lite
https://tensorflow.google.cn/lite/performance/quantization_spec

QAT 训练的模型具有与浮点相当的准确率

上表中，通过 TensorFlow Lite 默认配置训练得出 QAT 准确率，并与浮点基准和训练后量化模型进行了对比。

模拟低精度计算

训练计算图本身在浮点（例如 float32）中运行，但必须模拟定点的低精度计算（例如 TensorFlow Lite 一例中的 int8）。为此，我们在计算图中插入了特殊操作(tensorflow::ops::FakeQuantWithMinMaxVars)，以便将浮点张量转换为低精度，然后将低精度值转换回浮点值。这样可确保在计算中引入量化导致的损失，并确保后续计算模拟低精度计算。为实现这一目的，我们确保在张量中引入量化导致的损失，而由于浮点张量中的每个值现在都按 1:1 映射为低精度值，因此后续使用类似映射张量的所有计算都不会引入更多损失，并且将准确模拟低精度计算。

tensorflow::ops::FakeQuantWithMinMaxVars
https://tensorflow.google.cn/api_docs/cc/class/tensorflow/ops/fake-quant-with-min-max-vars

执行量化模拟操作

我们需要在训练计算图中执行量化模拟操作，使其符合量化计算图的计算方式。这表示，我们需要严格遵循 TensorFlow Lite 的量化规范，以确保 API 能够在 TensorFlow Lite 中执行。

量化规范
https://tensorflow.google.cn/lite/performance/quantization_spec

“‘wt quant”和“act quant”算子在模型的正向传递中引入了损失，以便模拟推理过程中的实际量化损失。请注意，Conv 和 ReLU6 之间没有量化操作。这是因为 ReLU 已融入 TensorFlow Lite

建立在 Keras 层和模型抽象层上的 API 隐藏了上述复杂操作，因此，您只需几行代码即可量化整个模型。

记录计算统计数据

除了模拟较低精度的计算，此 API 还负责记录量化训练模型所必需的统计数据。例如，您可以使用此 API 训练模型，并将其转换为只使用整数的 TensorFlow Lite 量化模型。

通过几行代码使用此 API

QAT API 为量化 TensorFlow Keras 模型提供了一种简单且高度灵活的方法。借此，您可以非常轻松地使用 “训练时量化” 训练整个或部分模型，然后将其导出，以便通过 TensorFlow Lite 进行部署。

量化整个 Keras 模型

import tensorflow_model_optimization as tfmot


model = tf.keras.Sequential([
   ...
])
# Quantize the entire model.
quantized_model = tfmot.quantization.keras.quantize_model(model)


# Continue with training as usual.
quantized_model.compile(...)
quantized_model.fit(...)

量化部分 Keras 模型

import tensorflow_model_optimization as tfmot
quantize_annotate_layer = tfmot.quantization.keras.quantize_annotate_layer


model = tf.keras.Sequential([
   ...
   # Only annotated layers will be quantized.
   quantize_annotate_layer(Conv2D()),
   quantize_annotate_layer(ReLU()),
   Dense(),
   ...
])


# Quantize the model.
quantized_model = tfmot.quantization.keras.quantize_apply(model)

默认情况下，我们将 API 配置为与 TensorFlow Lite 提供的量化执行支持一起使用。如需包含端到端训练示例的详细 Colab，请查看此处。

此处
https://tensorflow.google.cn/model_optimization/guide/quantization/training_example

此 API 非常灵活，并且能够处理更为复杂的用例。例如，您可以使用此 API 精确控制层内的量化、创建自定义量化算法，以及处理您编写的任何自定义层。

如需深入了解如何使用此 API，请试用此 Colab。Colab 中的这些部分提供了一些示例，展示了用户如何使用此 API 尝试不同的量化算法。您也可以观看近期 TensorFlow 开发者峰会的介绍。

此 Colab
https://tensorflow.google.cn/model_optimization/guide/quantization/training_comprehensive_guide
部分
https://tensorflow.google.cn/model_optimization/guide/quantization/training_comprehensive_guide#experiment_with_quantization

我们很高兴看到：QAT API 进一步帮助 TensorFlow 用户实现突破，在 TensorFlow Lite 支持的产品中实现更高效的执行，同时也为研究新的量化算法以及进一步开发具有不同精确率的硬件平台提供了机会。

如需了解详情，请观看下方的 TensorFlow 开发者峰会视频，该视频介绍并阐述了模型优化工具包和 QAT。

致谢

感谢 Pulkit Bhuwalka、Alan Chiao、Suharsh Sivakumar、Raziel Alvarez、Feng Liu、Lawrence Chan、Skirmantas Kligys、Yunlu Li、Khanh LeViet、Billy Lambert、Mark Daoust、Tim Davis、Sarah Sirajuddin 和 François Chollet。

了解更多请点击 “阅读原文” 访问官网。

— 推荐阅读 —