4. 量化感知训练入门

任务说明

本案例使用TensorFlow的量化感知训练（QAT）API，在CIFAR-10子集上训练一个小型CNN分类模型，完成以下闭环：

  1. 构建并训练一个带伪量化节点的基线模型；
  2. 导出为TensorFlow Lite格式（全整数量化）；
  3. 对比原始模型与量化模型的测试精度，评估量化损失。

环境依赖

需要 tensorflow==2.15.* 或兼容版本（含 tensorflow-model-optimization 库）。安装命令：

pip install tensorflow tensorflow-model-optimization

完整实现代码

以下代码完整实现了从数据加载、QAT训练到TFLite导出与精度对比的全部流程：

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

# 1. 加载CIFAR-10数据并降采样到32x32（原始尺寸），仅取前5000个训练样本和1000个测试样本
(x_train, y_train), (x_test, y_test) = tf.keras.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. 构建基础模型（不带量化）
def build_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(16, (3, 3), activation='relu', input_shape=(32, 32, 3)),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    return model

# 3. 应用量化感知训练（QAT）
qat_model = tfmot.quantization.keras.quantize_model(build_model())
qat_model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

# 4. 训练（仅2个epoch用于演示）
qat_model.fit(x_train, y_train, batch_size=128, epochs=2, validation_split=0.1, verbose=1)

# 5. 评估原始QAT模型精度
_, qat_acc = qat_model.evaluate(x_test, y_test, verbose=0)
print(f"QAT模型测试精度: {qat_acc:.4f}")

# 6. 转换为TFLite（全整数量化）
converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = lambda: ([x_train[i:i+1] for i in range(100)])
tflite_model = converter.convert()

# 7. 保存TFLite模型
with open('qat_cifar10.tflite', 'wb') as f:
    f.write(tflite_model)

# 8. 加载TFLite模型并评估精度
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

correct = 0
for i in range(len(x_test)):
    input_data = np.expand_dims(x_test[i], axis=0).astype(np.float32)
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    pred = np.argmax(interpreter.get_tensor(output_details[0]['index']))
    if pred == y_test[i]:
        correct += 1

tflite_acc = correct / len(x_test)
print(f"TFLite全整数量化模型测试精度: {tflite_acc:.4f}")
print(f"量化精度损失: {qat_acc - tflite_acc:.4f}")

运行方式

将上述代码保存为 qat_demo.py，在Python环境中直接执行：

python qat_demo.py

首次运行会自动下载CIFAR-10数据集（约170MB），需保持网络畅通。训练在CPU上约需2～3分钟，GPU环境会更快。

验证方法与通过标准

执行完成后，观察终端输出的三个关键数值：

  • QAT模型测试精度：应明显高于随机猜测（10%），通常在60%～75%区间（取决于训练随机性）；
  • TFLite全整数量化模型测试精度：应与QAT模型精度接近；
  • 量化精度损失：值应小于5个百分点，说明量化感知训练有效控制了精度下降。

同时确认当前目录下生成了 qat_cifar10.tflite 文件（约100～200KB），即代表转换成功。若精度损失超过5个百分点，可增加训练epoch数至5～10个重新运行。

常见问题

  • 数据集下载失败：检查网络或手动下载CIFAR-10数据集放置到 ~/.keras/datasets/ 目录；
  • 版本兼容错误：确保 tensorflow-model-optimization 版本与TensorFlow匹配，建议使用2.15.x；
  • 量化后精度大幅下降：增加训练轮次或使用更大训练子集（如10000样本）可改善。
← 上一章 📖 返回目录 下一章 →