📚 课程目录第2章 TFLite模型量化基础

2. TFLite模型量化基础

任务目标

本案例将带领你完成一个完整的TFLite模型量化流程：从一个训练好的浮点Keras模型出发，使用代表性数据集进行校准，最终转换为int8全量化模型。通过本案例，你将理解量化的工作原理，并掌握representative_dataset在量化过程中的核心作用。

环境依赖

TensorFlow 2.x（包含TFLiteConverter）
NumPy
Python 3.7+

完整实现

以下代码构建一个简单的浮点分类模型，训练后转换为int8量化模型。核心步骤包括：构建模型、训练、定义校准数据集、执行量化转换。

import numpy as np
import tensorflow as tf

# 1. 构建并训练一个简单的浮点模型（MNIST二分类演示）
def build_and_train_model():
    # 生成模拟数据：两类线性可分数据
    np.random.seed(42)
    x_train = np.concatenate([
        np.random.normal(2.0, 0.5, (500, 4)),
        np.random.normal(-2.0, 0.5, (500, 4))
    ], axis=0).astype(np.float32)
    y_train = np.array([0]*500 + [1]*500, dtype=np.int64)

    model = tf.keras.Sequential([
        tf.keras.layers.Dense(8, activation='relu', input_shape=(4,)),
        tf.keras.layers.Dense(2, activation='softmax')
    ])
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    model.fit(x_train, y_train, epochs=5, verbose=0)
    return model, x_train

# 2. 定义代表性数据集（校准用）
def representative_dataset(x_train, batch_size=32):
    # 从训练数据中采样，用于确定量化范围
    for i in range(0, len(x_train), batch_size):
        batch = x_train[i:i+batch_size]
        yield [batch]

# 3. 执行量化转换
def convert_to_int8(model, x_train):
    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
    return converter.convert()

# 4. 验证量化模型
def verify_quantized_model(tflite_model, x_train):
    interpreter = tf.lite.Interpreter(model_content=tflite_model)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    # 检查输入是否为int8
    assert input_details[0]['dtype'] == np.int8, "输入类型不是int8"
    assert output_details[0]['dtype'] == np.int8, "输出类型不是int8"

    # 量化参数
    input_scale, input_zero_point = input_details[0]['quantization']
    print(f"输入量化参数: scale={input_scale}, zero_point={input_zero_point}")

    # 用前10个样本测试
    test_samples = x_train[:10]
    for sample in test_samples:
        # 浮点输入量化为int8
        quantized_input = (sample / input_scale + input_zero_point).astype(np.int8)
        interpreter.set_tensor(input_details[0]['index'], [quantized_input])
        interpreter.invoke()
        quantized_output = interpreter.get_tensor(output_details[0]['index'])
        print(f"量化输出: {quantized_output[0]}")

    # 检查模型大小
    float_size = len(model.to_json()) + sum(w.numpy().nbytes for w in model.weights)
    int8_size = len(tflite_model)
    print(f"浮点模型大小约: {float_size} bytes")
    print(f"int8模型大小: {int8_size} bytes")

# 执行主流程
model, x_train = build_and_train_model()
tflite_int8_model = convert_to_int8(model, x_train)
verify_quantized_model(tflite_int8_model, x_train)

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

运行方式

安装依赖：pip install tensorflow numpy
将上述代码保存为quantize_demo.py
执行：python quantize_demo.py

验证方法与判定标准

类型验证：程序断言输入输出张量类型为int8，若断言通过则说明量化成功。
量化参数输出：观察打印的scale和zero_point，它们应为非零有限值。
推理输出：前10个样本的量化输出应显示为int8数值（0或1附近），表示模型仍能正确分类。
模型大小对比：int8模型文件应明显小于浮点模型（通常小4倍左右）。

关键说明

representative_dataset是int8量化的核心：它提供真实数据分布供转换器统计每个激活张量的数值范围，从而确定量化参数。若省略该函数，转换器只能进行动态范围量化（权重int8、激活float），无法实现全int8推理。本案例使用模拟数据训练模型，仅用于演示量化流程，实际部署时应使用真实任务数据。