TensorFlow 模型优化团队发布

我们很高兴在 TensorFlow 模型优化工具包 中发布量化感知训练 (QAT) API。QAT 使您能够训练和部署具有量化性能和大小优势的模型，同时保持接近其原始准确性。这项工作是我们 路线图 的一部分，旨在支持开发更小、更快的 ML 模型。有关更多背景信息，您可以查看之前关于 训练后量化、float16 量化 和 稀疏性 的文章。

量化是有损的

量化是将 ML 模型转换为使用较低精度参数和计算的等效表示的过程。这提高了模型的执行性能和效率。例如，TensorFlow Lite 8 位整数量化会导致模型的大小缩小 4 倍，计算速度提高 1.5 倍至 4 倍，并且 CPU 上的功耗降低。此外，它允许在专用神经加速器（例如 Coral 中的 Edge TPU）上执行模型，这些加速器通常有一套受限的数据类型。

但是，从更高精度到更低精度的过程本质上是有损的。如下图所示，量化将一小范围的浮点值压缩到固定数量的信息桶中。

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

这会导致信息丢失。模型的参数（或权重）现在只能取一小部分值，它们之间的微小差异会丢失。例如，范围 [2.0, 2.3] 中的所有值现在可能在同一个桶中表示。这类似于将小数表示为整数时出现的舍入误差。

还存在其他损失来源。当这些有损数字在多个乘加运算中使用时，这些损失会累积。此外，int8 值累积到 int32 整数后，需要重新缩放到 int8 值以进行下一次计算，从而引入更多计算误差。

量化感知训练

核心思想是在训练过程的前向传递中模拟低精度推理时间计算。这 项工作 归功于 Google Mobile Vision 团队的 Skirmantas Kligys 的原始创新。这将量化误差作为噪声引入训练过程，并作为整体损失的一部分，优化算法会尝试将其最小化。因此，模型会学习对量化更稳健的参数。

如果训练不可行，请查看 训练后量化，它作为 TensorFlow Lite 模型转换的一部分工作。QAT 对于希望尝试各种量化策略（超出 TensorFlow Lite 支持的范围）的研究人员和硬件设计师也非常有用，或者模拟量化如何影响不同硬件后端的准确性。

经过 QAT 训练的模型具有与浮点模型相当的准确性

在上表中，QAT 准确性数字使用默认 TensorFlow Lite 配置进行训练，并与浮点基线和训练后量化模型进行对比。

模拟低精度计算

训练图本身在浮点（例如 float32）中运行，但它必须模拟低精度计算，这是定点（例如 TensorFlow Lite 中的 int8）。为此，我们在图中插入特殊操作（tensorflow::ops::FakeQuantWithMinMaxVars），这些操作将浮点张量转换为低精度值，然后将低精度值转换回浮点。这确保量化中的损失在计算中被引入，并且进一步的计算模拟低精度。为此，我们确保量化带来的损失被引入张量中，并且由于浮点张量中的每个值现在都与一个低精度值一一对应，因此任何使用类似映射张量的进一步计算都不会引入任何进一步的损失，并且完全模仿低精度计算。

放置量化模拟操作

量化模拟操作需要放置在训练图中，使其与量化图的计算方式一致。这意味着，为了使我们的 API 能够在 TensorFlow Lite 中执行，我们需要精确地遵循 TensorFlow Lite 量化规范。

‘wt quant’ 和 ‘act quant’ 操作在模型的前向传递中引入损失，以模拟推理期间的实际量化损失。请注意，Conv 和 ReLU6 之间没有 Quant 操作。这是因为 ReLU 在 TensorFlow Lite 中被融合了。

该 API 基于 Keras 层和模型抽象构建，隐藏了上述复杂性，因此您可以用几行代码对整个模型进行量化。

记录计算统计信息

除了模拟降低精度的计算外，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 位于 这里。

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

要详细了解如何使用 API，请尝试使用 这个 Colab。Colab 中的这些 部分 提供了用户如何使用 API 尝试不同量化算法的示例。您还可以查看 TensorFlow 开发者峰会上的这个最近的演讲。

我们非常高兴看到 QAT API 如何进一步使 TensorFlow 用户能够突破其 TensorFlow Lite 驱动的产品中的高效执行的界限，以及它如何为研究新的量化算法和使用不同精度级别的进一步开发新硬件平台打开大门。

如果您想了解更多信息，请查看 TensorFlow DevSummit 中的这个视频，该视频介绍了模型优化工具包并解释了 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