训练后量化

训练后量化是一种转换技术,可以缩减模型大小,同时缩短 CPU 和硬件加速器的延迟时间,同时模型准确率几乎没有下降。使用 TensorFlow Lite Converter 将已经过训练的浮点 TensorFlow 模型转换为 TensorFlow Lite 格式时,您可以量化该模型。

优化方法

有多个训练后量化选项可供选择。下表总结了各种选项及其带来的益处:

方法 优势 硬件
动态范围量化 更小 4 倍,速度提升 2-3 倍 CPU
全整数量化 更小 4 倍,速度提升 3 倍 CPU、Edge TPU、微控制器
Float16 量化 缩小 2 倍,GPU 加速 CPU、GPU

以下决策树有助于确定哪种训练后量化方法最适合您的用例:

训练后优化选项

不量化

建议从转换为没有量化的 TFLite 模型开始。这将生成一个浮动 TFLite 模型。

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_quant_model = converter.convert()

我们建议您首先执行此操作,以验证原始 TF 模型的运算符是否与 TFLite 兼容,也可以将其用作基准来调试后续训练后量化方法引入的量化误差。例如,如果量化 TFLite 模型产生了意外结果,而浮点 TFLite 模型却准确,那么我们可以将问题范围缩小为由 TFLite 运算符的量化版本引入的错误。

动态范围量化

动态范围量化可减少内存使用量并加快计算速度,而无需提供代表性数据集进行校准。这种类型的量化仅在转换时以静态方式量化从浮点数转换为整数的权重,从而提供 8 位的精度:

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

为了进一步缩短推理过程中的延迟时间,“动态范围”运算符会根据激活数到 8 位的范围动态量化激活函数,并使用 8 位权重和激活函数执行计算。这项优化提供了接近完全定点推理的延迟时间。不过,输出仍使用浮点数进行存储,因此动态范围运算的提升速度小于完整的定点数计算。

全整数量化

通过确保所有模型数学运算都经过整数量化,您可以进一步缩短延迟时间、减少峰值内存用量,以及与仅支持整数的硬件设备或加速器的兼容性。

对于全整数量化,您需要校准或估算范围,即(min, max)。与权重和偏差等常量张量不同,变量张量(例如模型输入、激活(中间层的输出)和模型输出)无法校准,除非我们运行几个推理周期。因此,转换器需要使用代表性数据集来校准它们。此数据集可以是训练数据或验证数据的小型子集(大约 100-500 个样本)。请参阅下面的 representative_dataset() 函数。

在 TensorFlow 2.7 版本中,您可以通过签名来指定代表性数据集,如以下示例所示:

def representative_dataset():
  for data in dataset:
    yield {
      "image": data.image,
      "bias": data.bias,
    }

如果给定 TensorFlow 模型中有多个签名,您可以通过指定签名密钥来指定多个数据集:

def representative_dataset():
  # Feed data set for the "encode" signature.
  for data in encode_signature_dataset:
    yield (
      "encode", {
        "image": data.image,
        "bias": data.bias,
      }
    )

  # Feed data set for the "decode" signature.
  for data in decode_signature_dataset:
    yield (
      "decode", {
        "image": data.image,
        "hint": data.hint,
      },
    )

您可以通过提供输入张量列表来生成代表性数据集:

def representative_dataset():
  for data in tf.data.Dataset.from_tensor_slices((images)).batch(1).take(100):
    yield [tf.dtypes.cast(data, tf.float32)]

从 TensorFlow 2.7 版本开始,我们建议使用基于签名的方法而不是基于输入张量列表的方法,因为输入张量排序很容易翻转。

出于测试目的,您可以使用虚拟数据集,如下所示:

def representative_dataset():
    for _ in range(100):
      data = np.random.rand(1, 244, 244, 3)
      yield [data.astype(np.float32)]
 

具有浮点回退的整数(使用默认的浮点输入/输出)

为了对模型进行完全整数量化,但在浮点运算符没有整数实现时使用浮点运算符(以确保转换顺利进行),请按照以下步骤操作:

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
tflite_quant_model = converter.convert()

仅限整数

创建仅包含整数的模型是适用于微控制器的 TensorFlow LiteCoral Edge TPU 的常见用例。

此外,为了确保与纯整数设备(例如 8 位微控制器)和加速器(例如 Coral Edge TPU)兼容,您可以按照以下步骤对所有操作(包括输入和输出)强制执行全整数量化:

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8  # or tf.uint8
converter.inference_output_type = tf.int8  # or tf.uint8
tflite_quant_model = converter.convert()

Float16 量化

您可以通过将权重量化为 float16(16 位浮点数的 IEEE 标准)来减小浮点模型的大小。如需启用权重的 float16 量化,请按以下步骤操作:

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_quant_model = converter.convert()

float16 量化的优点如下:

  • 它可以将模型大小减少至原来的一半(因为所有权重都减至原始大小的一半)。
  • 它可以最大限度降低准确率损失。
  • 它支持一些可以直接对 float16 数据执行操作的委托(例如 GPU 委托),因此执行速度比 float32 计算更快。

float16 量化的缺点如下:

  • 延迟时间的减少幅度不及量化到定点数数学。
  • 默认情况下,在 CPU 上运行时,float16 量化模型会将权重值“反量化”为 float32。(请注意,GPU 代理不会执行此去量化,因为它可以对 float16 数据执行操作。)

仅限整数:具有 8 位权重的 16 位激活(实验性)

这是一个实验性的量化方案。它与“仅限整数”方案类似,但激活是根据其范围量化到 16 位,权重以 8 位整数量化,而偏差则量化为 64 位整数。这进一步称为 16x8 量化。

这种量化的主要优势在于,它可以显著提高准确率,但只会略微增加模型大小。

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.representative_dataset = representative_dataset
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8]
tflite_quant_model = converter.convert()

如果模型中的某些运算符不支持 16x8 量化,则模型仍可以量化,但不受支持的运算符将保持浮点数。为此,应将以下选项添加到 target_spec。

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.representative_dataset = representative_dataset
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
tf.lite.OpsSet.TFLITE_BUILTINS]
tflite_quant_model = converter.convert()

此量化方案可提高准确性的用例示例包括:

  • super-resolution,
  • 音频信号处理,例如降噪和波束成形;
  • 图像去噪、
  • 根据单张图片进行 HDR 重构。

这种量化的缺点是:

  • 由于缺少优化的内核实现,目前的推理速度明显慢于 8 位全整数。
  • 目前,它与现有的硬件加速 TFLite 代理不兼容。

您可以在此处找到关于此量化模式的教程。

模型准确率

由于权重在训练后进行了量化,因此准确性可能会降低,尤其是对于较小的网络。Kaggle 模型上为特定网络提供预训练的完全量化模型。请务必检查量化模型的准确率,以验证准确性的任何下降是否在可接受的范围内。我们提供了可用于评估 TensorFlow Lite 模型准确率的工具。

或者,如果准确率下降太高,请考虑使用量化感知训练。但是,这样做需要在模型训练期间进行修改,以添加虚假量化节点,而本页面上的训练后量化技术使用的是现有的预训练模型。

量化张量的表示法

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

\[real\_value = (int8\_value - zero\_point) \times scale\]

该表示法包含两个主要部分:

  • 每个轴(也称为每个通道)的权重或每个张量权重,以 int8 二的补码值表示,范围为 [-127, 127],零点等于 0。

  • 每个张量激活/输入以 int8 二补数表示,范围为 [-128, 127],零点的范围为 [-128, 127]。

如需详细了解量化方案,请参阅我们的量化规范。我们鼓励希望插入 TensorFlow Lite 委托接口的硬件供应商实现此处所述的量化方案。