4. TensorFlow模型训练与导出：Keras模型构建、训练后量化（PTQ）、导出为TFLite格式

这一章，我们终于要动手把模型跑起来了。说实话，很多做MCU部署的朋友，卡就卡在模型导出这一步——训练好好的，一量化就崩，一转TFLite就报错。别急，我带你一步步走通这个流程。

4.1 用Keras搭一个能上MCU的模型

先说说我的选型原则。MCU上跑模型，资源就那么点——Flash可能就512KB，RAM可能就256KB。所以模型不能太胖。我个人习惯用深度可分离卷积，参数量只有普通卷积的1/8到1/10。

下面这个例子，是一个针对CIFAR-10的轻量级模型。你直接拿去改改就能用。

import tensorflow as tf
from tensorflow.keras import layers, models

def build_lightweight_model(input_shape=(32, 32, 3), num_classes=10):
    model = models.Sequential([
        # 第一层：普通卷积，提取基础特征
        layers.Conv2D(16, (3, 3), padding='same', input_shape=input_shape),
        layers.BatchNormalization(),
        layers.ReLU(),
        layers.MaxPooling2D((2, 2)),

        # 第二层：深度可分离卷积，省参数
        layers.SeparableConv2D(32, (3, 3), padding='same'),
        layers.BatchNormalization(),
        layers.ReLU(),
        layers.MaxPooling2D((2, 2)),

        # 第三层：继续压缩
        layers.SeparableConv2D(64, (3, 3), padding='same'),
        layers.BatchNormalization(),
        layers.ReLU(),
        layers.GlobalAveragePooling2D(),

        # 分类头
        layers.Dense(num_classes, activation='softmax')
    ])
    return model

model = build_lightweight_model()
model.summary()

4.2 训练时就要为量化做准备

很多人训练完模型直接扔给量化工具，结果惨不忍睹。为什么？因为训练时没考虑量化误差。

我建议你在训练时做两件事：

1. 使用伪量化训练（QAT）——虽然这章讲的是PTQ，但训练时加一点量化噪声，能让模型对量化更鲁棒。
2. 监控激活值范围——如果某个层的输出范围特别大（比如超过[-6, 6]），量化精度损失就会很严重。

下面是一个简单的训练脚本，加了学习率衰减和早停：

# 数据准备
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# 编译
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 回调函数
callbacks = [
    tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3),
    tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
]

# 训练
history = model.fit(
    x_train, y_train,
    batch_size=64,
    epochs=50,
    validation_data=(x_test, y_test),
    callbacks=callbacks
)

# 保存模型
model.save('cifar10_lightweight.h5')

4.3 训练后量化（PTQ）——最省事的量化方式

PTQ，说白了就是模型训练完了，我再给它做一次"瘦身"。不需要重新训练，只需要一小部分校准数据。

它的核心流程是这样的：

• 收集一小批代表性数据（一般100-500张图就够了）
• 跑一遍推理，统计每层的激活值范围
• 根据统计结果，计算量化参数（scale和zero_point）
• 把权重和激活值从float32映射到int8

我画了一张流程图，帮你理清PTQ的完整链路：

4.4 动手做PTQ：代码实战

好，理论说完了，咱们直接上代码。TensorFlow提供了tf.lite.TFLiteConverter，一行代码就能完成量化。

import tensorflow as tf

# 加载训练好的模型
model = tf.keras.models.load_model('cifar10_lightweight.h5')

# 准备校准数据（取训练集前200张）
calibration_data = x_train[:200]

# 定义校准数据生成器
def representative_dataset():
    for i in range(200):
        # 每次返回一个batch，shape为(1, 32, 32, 3)
        yield [calibration_data[i:i+1].astype(np.float32)]

# 转换并量化
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

tflite_quant_model = converter.convert()

# 保存量化后的模型
with open('cifar10_quant_int8.tflite', 'wb') as f:
    f.write(tflite_quant_model)

print(f"量化模型大小：{len(tflite_quant_model) / 1024:.2f} KB")

4.5 量化前后对比：数据说话

我拿CIFAR-10测试集跑了一下，结果如下。你想想看，模型体积缩小了4倍，推理速度提升3-4倍，精度只掉了不到1%。这笔账怎么算都划算。

4.6 导出TFLite后的验证

模型导出完了，别急着部署。先在本机验证一下，确保量化后的模型推理结果正确。

import numpy as np

# 加载量化模型
interpreter = tf.lite.Interpreter(model_path='cifar10_quant_int8.tflite')
interpreter.allocate_tensors()

# 获取输入输出信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print(f"输入类型：{input_details[0]['dtype']}")
print(f"输入量化参数：scale={input_details[0]['quantization'][0]}, "
      f"zero_point={input_details[0]['quantization'][1]}")

# 准备测试数据（需要量化到int8）
test_image = x_test[0:1]  # shape: (1, 32, 32, 3)
input_scale, input_zp = input_details[0]['quantization']
test_image_quant = (test_image / input_scale + input_zp).astype(np.int8)

# 推理
interpreter.set_tensor(input_details[0]['index'], test_image_quant)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])

# 反量化输出
output_scale, output_zp = output_details[0]['quantization']
output_float = (output.astype(np.float32) - output_zp) * output_scale

print(f"预测结果：{np.argmax(output_float)}")
print(f"真实标签：{y_test[0]}")