TensorFlow 模型优化团队发布

自从我们发布了模型优化工具包——一套既可供新手开发者也可供高级开发者使用的优化机器学习模型部署和执行的技术——我们一直在努力简化量化机器学习模型的复杂性。

最初，我们支持“混合操作”方式的训练后量化，也就是量化模型参数（即权重），但允许部分计算在浮点中进行。今天，我们很高兴宣布我们的工具的下一个新增功能：训练后整数量化。整数量化是一种通用技术，它可以降低模型权重和激活值的数值精度，从而减小内存占用并提高延迟。

量化模型以减小尺寸、延迟和功耗，同时保持几乎不损失精度

为什么要使用训练后整数量化

我们之前发布的“混合”训练后量化方法在许多情况下都减小了模型尺寸和延迟，但它有一个限制：需要浮点计算，而并非所有硬件加速器（例如，Edge TPU）都提供浮点计算，但这使得它适合于 CPU。

我们新的训练后整数量化允许用户将已经训练好的浮点模型完全量化为仅使用 8 位有符号整数（即 `int8`）。通过利用这种量化方案，我们可以在许多模型中获得合理的量化模型精度，而无需使用量化感知训练来重新训练模型。有了这个新工具，模型将继续缩小 4 倍，但将看到更大的 CPU 加速。固定点硬件加速器，如 Edge TPU，也将能够运行这些模型。

与量化感知训练相比，此工具使用起来简单得多，并且在大多数模型上都提供了可比的精度。量化感知训练可能仍然有一些用例，但我们预计随着我们继续改进训练后工具，这种情况将变得越来越少。

总之，用户在针对简单的 CPU 尺寸和延迟改进时应该使用“混合”训练后量化。当针对更大的 CPU 改进或固定点加速器时，他们应该使用此整数训练后量化工具，如果模型的精度下降，则可以使用量化感知训练。

如何启用训练后整数量化

我们的整数量化工具需要一小部分代表性数据的校准集。只需将 `representative_dataset` 生成器提供给转换器，优化参数就会对输入模型执行整数量化。

def representative_dataset_gen():
  data = tfds.load(...)

  for _ in range(num_calibration_steps):
    image, = data.take(1)
    yield [image]

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = tf.lite.RepresentativeDataset(
    representative_dataset_gen) 

模型是否完全量化？

与现有的训练后量化功能一样，默认情况下，没有量化实现的运算 (“ops”) 会自动保留为浮点。这允许转换顺利进行，并会生成一个始终在典型移动 CPU 上执行的模型——考虑到 TensorFlow Lite 将在仅整数加速器中执行整数运算，并为涉及浮点的运算回退到 CPU。要完全在不支持浮点运算的专用硬件（例如，一些机器学习加速器，包括 Edge TPU）上执行，您可以指定一个标志以仅输出整数运算

converter.target_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

当使用此标志并且运算没有整数可量化对应物时，TensorFlow Lite 转换器将抛出错误。

只需少量数据

在我们的实验中，我们发现，几十个代表模型在执行期间将看到的示例就足以获得最佳精度。例如，下面的精度数字来自仅在 ImageNet 数据集的 100 张图像上进行校准的模型。

结果

延迟

与浮点模型相比，量化模型在 CPU 上快 2-4 倍，并且缩小 4 倍。我们预计硬件加速器（如 Edge TPU）将进一步加快速度。

精度

仅使用 ImageNet 数据集的 100 张校准图像，完全量化的整数模型的精度与其浮点版本相当（MobileNet v1 损失 1%）。

这些整数模型的工作原理

记录动态范围

我们的新工具通过记录动态范围来工作，在浮点 TensorFlow Lite 模型上运行多次推理，使用用户提供的代表性数据集作为输入。我们使用从推理中记录的值来确定执行模型中所有张量所需的整数运算的缩放参数。

Int8 量化方案

需要注意的是，我们的新量化规范使这种使用每轴量化的训练后用例成为可能。在添加每轴量化之前，训练后整数量化由于精度下降而变得不切实际；但每轴的精度优势使许多模型的精度更接近浮点。

8 位量化使用以下公式来近似浮点值

real_value = (sint8_value — zero_point) * scale。

每轴（也称为“每通道”）或每层权重由范围 [-127, 127] 中的 int8 二进制补码值表示，零点等于 0。

每层激活/输入由范围 [-128, 127] 中的 int8 二进制补码值表示，零点在范围 [-128, 127] 内。

有关更多详细信息，请参阅完整量化规范。

量化感知训练怎么样？

我们相信让量化尽可能简单。因此，能够在训练后量化模型是我们非常兴奋的事情！但是，我们也知道，有些模型在用量化训练时才能保持最佳质量。这就是我们也在开发量化感知训练 API 的原因。与此同时，我们鼓励您尝试训练后量化，因为这可能是您的模型所需要的全部！

文档和教程

在 TensorFlow 网站上，您可以找到有关训练后整数量化、我们的新量化规范以及训练后整数量化教程的更多信息。我们很乐意听取您使用此功能的体验——分享您的故事！

鸣谢

Suharsh Sivakumar、Jian Li、Shashi Shekhar、Yunlu Li、Alan Chiao、Raziel Alvarez、Lawrence Chan、Daniel Situnayake、Tim Davis、Sarah Sirajuddin