📚 课程目录第4章使用TensorFlow Lite对MobileNetV2进行后训练整数量化

4. 使用TensorFlow Lite对MobileNetV2进行后训练整数量化

任务目标

将预训练的 MobileNetV2 模型转换为 TensorFlow Lite 格式，应用后训练整数量化（PTQ），并在 CPU 上对比量化前后模型的推理速度与精度变化。本案例聚焦于量化流程的完整闭环，不涉及边缘设备的具体部署。

环境依赖

TensorFlow 2.x（含 tensorflow-model-optimization 可选，本案例使用原生 TFLiteConverter 即可）
NumPy、Pillow（用于图像预处理）
ImageNet 验证集子集（或任意 100 张测试图片，需与标签对应）

完整实现代码

import tensorflow as tf
import numpy as np
import pathlib
import time

# 1. 加载预训练模型并转换
model = tf.keras.applications.MobileNetV2(weights='imagenet', input_shape=(224, 224, 3))
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# 2. 后训练整数量化配置
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset_gen  # 在下方定义

# 3. 转换并保存
tflite_quant_model = converter.convert()
with open('mobilenetv2_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)

# 4. 加载量化模型进行推理
interpreter = tf.lite.Interpreter(model_path='mobilenetv2_quant.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 5. 精度与速度评估函数
def evaluate_tflite(interpreter, images, labels, num_runs=10):
    """返回 top-1 精度和平均推理延迟(ms)"""
    correct = 0
    latencies = []
    for i, img in enumerate(images):
        # 预处理：归一化到 [-1, 1]
        img = (img.astype(np.float32) / 127.5) - 1.0
        img = np.expand_dims(img, axis=0)

        # 推理计时（前 5 次预热）
        if i < 5:
            interpreter.set_tensor(input_details[0]['index'], img)
            interpreter.invoke()
            continue

        start = time.perf_counter()
        interpreter.set_tensor(input_details[0]['index'], img)
        interpreter.invoke()
        latencies.append((time.perf_counter() - start) * 1000)

        output = interpreter.get_tensor(output_details[0]['index'])
        pred = np.argmax(output[0])
        if pred == labels[i]:
            correct += 1

    avg_latency = np.mean(latencies) if latencies else 0
    accuracy = correct / max(1, len(images) - 5)  # 去掉预热样本
    return accuracy, avg_latency

# 6. 数据加载（示例：从目录读取）
def load_images(data_dir, num_samples=100):
    """从 data_dir 读取图片，返回 (images, labels)"""
    data_dir = pathlib.Path(data_dir)
    image_paths = list(data_dir.glob('*.jpg'))[:num_samples]
    images, labels = [], []
    for path in image_paths:
        img = tf.keras.preprocessing.image.load_img(path, target_size=(224, 224))
        img = tf.keras.preprocessing.image.img_to_array(img)
        images.append(img)
        # 假设文件名格式为 "label_xxx.jpg"
        label = int(path.stem.split('_')[0])
        labels.append(label)
    return np.array(images), np.array(labels)

# 7. 代表性数据集生成器（用于量化校准）
def representative_dataset_gen():
    # 使用训练集中的 100 张图片作为校准数据
    calib_images, _ = load_images('calib_images', num_samples=100)
    for img in calib_images:
        img = (img.astype(np.float32) / 127.5) - 1.0
        yield [np.expand_dims(img, axis=0)]

# 8. 主流程
if __name__ == '__main__':
    # 准备测试数据（需提前准备）
    test_images, test_labels = load_images('test_images', num_samples=100)

    # 评估量化模型
    quant_acc, quant_lat = evaluate_tflite(interpreter, test_images, test_labels)

    # 评估原始 Keras 模型（作为对比基准）
    def evaluate_keras(model, images, labels):
        correct = 0
        latencies = []
        for i, img in enumerate(images):
            img = (img.astype(np.float32) / 127.5) - 1.0
            img = np.expand_dims(img, axis=0)
            if i < 5:
                model.predict(img, verbose=0)
                continue
            start = time.perf_counter()
            pred = model.predict(img, verbose=0)
            latencies.append((time.perf_counter() - start) * 1000)
            if np.argmax(pred[0]) == labels[i]:
                correct += 1
        return correct / max(1, len(images) - 5), np.mean(latencies)

    keras_acc, keras_lat = evaluate_keras(model, test_images, test_labels)

    print(f'Keras 模型   - 精度: {keras_acc:.4f}, 延迟: {keras_lat:.2f} ms')
    print(f'TFLite 量化  - 精度: {quant_acc:.4f}, 延迟: {quant_lat:.2f} ms')
    print(f'模型大小: {pathlib.Path("mobilenetv2_quant.tflite").stat().st_size / 1024:.1f} KB')

运行与验证

准备数据：创建 calib_images/ 和 test_images/ 目录，放入 ImageNet 子集图片，文件名格式为 {label}_{任意名}.jpg（label 为 0-999 的整数）。
执行脚本：python quantize_mobilenetv2.py
通过条件：
- 脚本正常输出量化后模型文件 mobilenetv2_quant.tflite；
- 量化模型精度与 Keras 模型相比下降不超过 2%；
- 量化模型推理延迟明显低于 Keras 模型（通常快 2-5 倍）；
- 量化模型文件大小约为原始模型的 1/4（约 4-5 MB）。

关键说明

代表性数据集需覆盖各类别分布，通常取 100-500 张即可获得良好量化效果；
若精度下降超过预期，可尝试增加代表性数据集样本量；
本案例使用 CPU 推理延迟作为速度指标，实际边缘设备（如树莓派、手机）上的加速比会更高。

上述流程展示了从 Keras 模型到 TFLite INT8 量化模型的完整转换链路。核心步骤在于 representative_dataset_gen 提供的校准数据，它决定了量化时激活值的动态范围，直接影响量化精度。运行脚本后，观察输出的精度与延迟数据，若满足通过条件即完成本案例目标。