3. TensorFlow Lite模型转换与INT8量化

任务目标

将训练好的TensorFlow浮点模型转换为TensorFlow Lite格式，并应用INT8全整数量化，使模型权重与激活均以8位整数表示，以适配低算力嵌入式设备。转换完成后，通过对比量化前后模型在验证集上的推理输出，确认量化误差在可接受范围内。

环境与依赖

需要TensorFlow 2.x（含TFLite Converter），以及用于图像预处理的Pillow和NumPy。若需在PC上验证推理一致性，还需安装TFLite Runtime（TensorFlow 2.x自带）。

完整实现

以下代码完成：构建并训练一个简单的MNIST分类模型（仅用于演示量化流程），然后执行INT8全整数量化转换，最后加载量化模型对单张测试图推理并输出量化前后对比。

import tensorflow as tf
import numpy as np
from PIL import Image

# 1. 构建并训练一个简单模型（演示用）
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(28, 28, 1)),
    tf.keras.layers.Conv2D(8, 3, activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10)
])
model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])
model.fit(x_train[..., tf.newaxis], y_train, epochs=1, batch_size=64, validation_split=0.1, verbose=0)

# 2. 定义代表验证集的校准数据集（取100张）
calib_ds = tf.data.Dataset.from_tensor_slices((x_test[:100, ..., tf.newaxis].astype(np.float32), y_test[:100]))
def representative_dataset():
    for data, _ in calib_ds.batch(1).take(100):
        yield [data]

# 3. 转换为TFLite并应用INT8全整数量化
converter = tf.lite.TFLiteConverter.from_keras_model(model)
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
converter.inference_output_type = tf.int8
tflite_int8_model = converter.convert()

# 保存量化模型
with open('model_int8.tflite', 'wb') as f:
    f.write(tflite_int8_model)

# 4. 加载量化模型并推理单张图片
interpreter = tf.lite.Interpreter(model_content=tflite_int8_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 输入预处理：归一化到[0,1]后按量化参数缩放
img = x_test[0, ..., tf.newaxis].astype(np.float32)
input_scale, input_zero = input_details[0]['quantization']
img_int8 = (img / input_scale + input_zero).astype(np.int8)
interpreter.set_tensor(input_details[0]['index'], img_int8[None, ...])
interpreter.invoke()
output_int8 = interpreter.get_tensor(output_details[0]['index'])

# 反量化输出
out_scale, out_zero = output_details[0]['quantization']
output_float = (output_int8.astype(np.float32) - out_zero) * out_scale

# 5. 对比浮点模型输出
float_logits = model.predict(x_test[0, ..., tf.newaxis][None, ...], verbose=0)[0]
print("浮点模型预测类别:", np.argmax(float_logits))
print("INT8量化模型预测类别:", np.argmax(output_float))
print("真实标签:", y_test[0])
print("量化前后logits差异(L2范数):", np.linalg.norm(float_logits - output_float))

运行与验证

直接执行上述脚本即可完成全流程。运行前确保已安装TensorFlow 2.x。脚本将输出量化前后模型的预测类别及logits差异值。

通过条件：量化前后模型对同一测试样本的预测类别一致，且logits的L2范数差异小于5.0（MNIST简单任务典型差异通常小于2.0）。若差异过大，需检查校准数据集是否具有代表性或增加校准样本数量。

关键说明

• INT8量化要求模型输入输出均为int8类型，因此推理前需按量化参数对输入进行缩放转换，推理后需反量化输出才能与浮点结果比较。
• 校准数据集仅用于统计各层激活值的动态范围，不参与训练，通常取100～500张验证集图片即可。
• 若目标设备不支持全整型运算（如部分GPU加速器），可移除target_spec.supported_ops限制，转换器将自动回退为float16或动态范围量化。