📚 课程目录第3章量化感知训练模拟

3. 量化感知训练模拟

任务目标

在TensorFlow/Keras中实现量化感知训练（QAT）模拟，通过在训练过程中插入伪量化节点，使模型权重与激活分布适应INT8量化误差，最终导出精度更高的INT8模型。本案例使用CIFAR-10子集训练一个小型CNN，对比普通训练与QAT训练后INT8模型的精度差异。

环境依赖

Python 3.8+
TensorFlow 2.12+（含TensorFlow Model Optimization Toolkit）
NumPy

安装命令：pip install tensorflow tensorflow-model-optimization numpy

关键完整代码

import tensorflow as tf
import tensorflow_model_optimization as tfmot
import numpy as np
from tensorflow.keras import layers, datasets

# 1. 加载CIFAR-10子集（仅用前5000张训练，1000张测试，加速演示）
(x_train, y_train), (x_test, y_test) = datasets.cifar10.load_data()
x_train, x_test = x_train[:5000] / 255.0, x_test[:1000] / 255.0
y_train, y_test = y_train[:5000].flatten(), y_test[:1000].flatten()

# 2. 定义基础CNN模型
def build_model():
    model = tf.keras.Sequential([
        layers.Conv2D(16, 3, activation='relu', input_shape=(32, 32, 3)),
        layers.MaxPooling2D(),
        layers.Conv2D(32, 3, activation='relu'),
        layers.MaxPooling2D(),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(10, activation='softmax')
    ])
    return model

# 3. 训练普通模型（基线）
base_model = build_model()
base_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
base_model.fit(x_train, y_train, epochs=3, batch_size=128, validation_split=0.1, verbose=0)

# 4. 应用QAT伪量化
qat_model = tfmot.quantization.keras.quantize_model(build_model())
qat_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
qat_model.fit(x_train, y_train, epochs=3, batch_size=128, validation_split=0.1, verbose=0)

# 5. 转换为INT8 TFLite模型
def convert_int8(model, name):
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = lambda: ([x_train[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
    tflite_model = converter.convert()
    with open(name, 'wb') as f:
        f.write(tflite_model)
    return tflite_model

base_int8 = convert_int8(base_model, 'base_int8.tflite')
qat_int8 = convert_int8(qat_model, 'qat_int8.tflite')

# 6. 评估INT8模型精度
def evaluate_tflite(model_path):
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    correct = 0
    for i in range(200):  # 用200张测试图快速验证
        input_data = x_test[i:i+1].astype(np.float32)
        # 输入量化
        scale, zero = input_details[0]['quantization']
        input_int8 = (input_data / scale + zero).astype(np.int8)
        interpreter.set_tensor(input_details[0]['index'], input_int8)
        interpreter.invoke()
        output = interpreter.get_tensor(output_details[0]['index'])
        # 输出反量化
        scale, zero = output_details[0]['quantization']
        output_float = (output.astype(np.float32) - zero) * scale
        if np.argmax(output_float) == y_test[i]:
            correct += 1
    return correct / 200

print(f"基线INT8精度: {evaluate_tflite('base_int8.tflite'):.3f}")
print(f"QAT INT8精度: {evaluate_tflite('qat_int8.tflite'):.3f}")

运行与验证

将上述代码保存为qat_demo.py并执行：python qat_demo.py
程序会训练两个模型（普通与QAT），分别转换为INT8 TFLite格式，并输出200张测试图上的精度。
通过条件：QAT模型的INT8精度应高于或等于基线模型的INT8精度（通常提升2~8个百分点）。若两者精度接近或QAT略低，可增加训练轮次至5~10轮再比较。

关键说明

第4步quantize_model()会在各层间插入伪量化节点，训练时模拟INT8舍入误差，使权重适应低精度表示。
第5步转换时representative_dataset提供100张校准图片，用于确定激活值的动态范围，是INT8量化必需步骤。
评估函数中手动处理了输入/输出的INT8量化与反量化，与TFLite运行时行为一致。
若需部署到边缘设备，生成的qat_int8.tflite可直接被TensorFlow Lite解释器加载。