📚 课程目录第3章 TFLite模型转换与量化感知训练

3. TFLite模型转换与量化感知训练

任务目标

本案例使用TensorFlow 2.x，完成以下闭环：构建一个简单的Keras图像分类模型 → 应用量化感知训练（QAT）微调 → 转换为TFLite格式 → 对比普通转换与QAT转换后的模型大小及推理输出。核心目标是验证QAT流程能显著减小TFLite模型体积，同时保持推理结果可用。

环境依赖

需要 tensorflow==2.13+（含 tensorflow-model-optimization）。安装命令：

pip install tensorflow tensorflow-model-optimization

完整实现代码

import tensorflow as tf
import tensorflow_model_optimization as tfmot
import numpy as np

# 1. 构建简单Keras模型（MNIST 10分类）
def build_model():
    return 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)
    ])

# 2. 准备少量示例数据（仅用于演示转换流程）
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train[:2000].astype('float32') / 255.0
x_train = x_train[..., np.newaxis]
y_train = tf.keras.utils.to_categorical(y_train[:2000], 10)

# 3. 训练基线模型（2个epoch即可）
model = build_model()
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=2, batch_size=64, verbose=0)

# 4. 应用QAT（只量化全模型）
qat_model = tfmot.quantization.keras.quantize_model(model)
qat_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
qat_model.fit(x_train, y_train, epochs=1, batch_size=64, verbose=0)

# 5. 转换函数：普通 vs QAT
def convert_tflite(model, use_qat=False):
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    if use_qat:
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
    tflite_model = converter.convert()
    return tflite_model

# 6. 执行转换并对比大小
baseline_tflite = convert_tflite(model, use_qat=False)
qat_tflite = convert_tflite(qat_model, use_qat=True)

print(f"普通TFLite大小: {len(baseline_tflite)} bytes")
print(f"QAT+TFLite大小: {len(qat_tflite)} bytes")
print(f"体积缩减: {100 * (1 - len(qat_tflite) / len(baseline_tflite)):.1f}%")

# 7. 验证推理一致性（使用同一输入）
def run_inference(tflite_model, input_data):
    interpreter = tf.lite.Interpreter(model_content=tflite_model)
    interpreter.allocate_tensors()
    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'])

test_input = x_train[0:1]
out_baseline = run_inference(baseline_tflite, test_input)
out_qat = run_inference(qat_tflite, test_input)

print("普通模型预测类别:", np.argmax(out_baseline))
print("QAT模型预测类别:", np.argmax(out_qat))
print("预测一致:", np.argmax(out_baseline) == np.argmax(out_qat))

运行方式

将上述代码保存为 qat_demo.py，直接执行：

python qat_demo.py

首次运行会自动下载MNIST数据集（约10MB），需要网络连接。

验证方法与通过标准

体积对比：控制台输出中，QAT+TFLite大小应明显小于普通TFLite大小，体积缩减通常在70%以上（因权重从float32量化为int8）。
推理一致性：最后一行输出预测一致: True，表明量化后模型对同一输入的预测类别与原始模型相同，说明QAT流程有效保留了模型决策能力。
无报错：脚本完整执行无异常，且所有输出数值均为实际运行产生。

关键说明

本案例使用MNIST前2000个样本、2个epoch训练，仅用于演示QAT转换流程，不追求高精度。实际项目中应使用完整数据集训练至收敛后再应用QAT。
QAT的核心机制是在训练中模拟量化误差（伪量化节点），使模型权重适应低比特表示，从而在转换后保持精度。本例中quantize_model自动在层间插入伪量化节点。
转换时QAT模型必须配合converter.optimizations = [tf.lite.Optimize.DEFAULT]才能产出int8量化模型；普通模型即使加此参数也不会自动量化（需额外校准数据），因此本例中普通模型使用float32转换作为对照基线。