📚 课程目录第3章模型量化进阶

3. 模型量化进阶

任务目标

本案例将带领你完成一个分类模型的全整数量化（Full Integer Quantization）流程：使用校准数据集确定量化参数，将模型权重与激活从浮点转换为 int8，并生成一个完全脱离浮点运算的 TensorFlow Lite 模型。最后，通过对比量化前后模型在相同输入下的输出，验证嵌入式设备上推理的数值一致性。

环境与依赖

Python 3.8+
TensorFlow 2.12+（包含 tflite-runtime 或完整版以支持量化与推理）
NumPy

完整实现

以下代码完成：构建并训练一个简单 CNN 模型（MNIST 手写数字识别）→ 转换为全整数量化模型 → 使用代表性数据集校准 → 验证数值一致性。

import tensorflow as tf
import numpy as np

# 1. 加载并预处理 MNIST 数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.astype(np.float32) / 255.0
x_test = x_test.astype(np.float32) / 255.0
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

# 2. 构建并训练一个简单 CNN（为演示量化流程，仅训练 1 个 epoch）
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(8, 3, activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D(2),
    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, y_train, epochs=1, batch_size=128, validation_split=0.1, verbose=0)

# 3. 转换为 TensorFlow Lite 模型（浮点基线）
converter = tf.lite.TFLiteConverter.from_keras_model(model)
float_lite_model = converter.convert()

# 4. 全整数量化转换
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = lambda: ([x_test[i:i+1] for i in range(100)])  # 校准数据集
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
int8_lite_model = converter.convert()

# 5. 保存模型文件
with open('mnist_float.tflite', 'wb') as f:
    f.write(float_lite_model)
with open('mnist_int8.tflite', 'wb') as f:
    f.write(int8_lite_model)

# 6. 验证数值一致性
def run_inference(interpreter, input_data):
    """对单个输入样本执行推理，返回输出张量"""
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    return interpreter.get_tensor(output_details[0]['index'])

# 6.1 浮点模型推理
float_interp = tf.lite.Interpreter(model_content=float_lite_model)
float_interp.allocate_tensors()
float_input_detail = float_interp.get_input_details()[0]
float_output_detail = float_interp.get_output_details()[0]

# 6.2 全整数量化模型推理
int8_interp = tf.lite.Interpreter(model_content=int8_lite_model)
int8_interp.allocate_tensors()
int8_input_detail = int8_interp.get_input_details()[0]
int8_output_detail = int8_interp.get_output_details()[0]

# 6.3 量化参数（用于将 float 输入转换为 int8）
input_scale, input_zero_point = int8_input_detail['quantization']
output_scale, output_zero_point = int8_output_detail['quantization']

# 6.4 选取 10 个测试样本进行对比
num_test = 10
max_diff = 0.0
for i in range(num_test):
    # 浮点输入
    float_input = x_test[i:i+1]
    
    # 浮点模型推理
    float_interp.set_tensor(float_input_detail['index'], float_input)
    float_interp.invoke()
    float_output = float_interp.get_tensor(float_output_detail['index']).flatten()
    
    # 将输入量化为 int8
    int8_input = (float_input / input_scale + input_zero_point).astype(np.int8)
    int8_interp.set_tensor(int8_input_detail['index'], int8_input)
    int8_interp.invoke()
    int8_output_raw = int8_interp.get_tensor(int8_output_detail['index']).flatten()
    
    # 反量化输出
    int8_output = (int8_output_raw.astype(np.float32) - output_zero_point) * output_scale
    
    # 计算最大差异
    diff = np.max(np.abs(float_output - int8_output))
    max_diff = max(max_diff, diff)
    
    # 检查 argmax 是否一致（预测类别）
    if np.argmax(float_output) != np.argmax(int8_output):
        print(f"样本 {i}: 类别不一致! float={np.argmax(float_output)}, int8={np.argmax(int8_output)}")

print(f"量化前后最大输出差异: {max_diff:.6f}")
print("验证完成。若最大差异小于 0.1 且类别一致，则数值一致性验证通过。")

运行与验证

运行脚本：将上述代码保存为 full_int8_quantization.py，在终端执行 python full_int8_quantization.py。
观察输出：脚本会打印量化前后输出的最大差异。若差异小于 0.1 且所有测试样本的预测类别一致，则验证通过。
模型文件：运行后会在当前目录生成 mnist_float.tflite（浮点基线）和 mnist_int8.tflite（全整数量化）两个文件，可直接部署到嵌入式设备。

关键说明

校准数据集：代码中使用测试集前 100 个样本作为代表性数据集（representative_dataset），用于统计激活值的动态范围，从而确定量化参数。实际应用中应使用训练集或验证集中有代表性的子集。
输入/输出类型：通过 inference_input_type 和 inference_output_type 强制将模型的输入输出也量化为 int8，确保模型完全脱离浮点运算。
数值一致性验证：由于量化会引入精度损失，输出差异通常在 0.01～0.1 范围内。若差异过大，可增加校准数据集规模或使用更大位宽（如 int16）的量化。