模型准备与格式转换：从训练到部署的第一步

说实话，模型训练只是万里长征走了一半。真正让模型在 STM32 上跑起来，中间还有一道绕不过去的坎——格式转换。我见过太多人训练时精度漂亮得很，一到部署就翻车，十有八九是格式转换这步没处理好。

这一章，咱们就聊聊怎么把 TensorFlow/Keras 模型变成 STM32Cube.AI 能吃的格式。说白了，就是给模型「换装」的过程。

为什么需要格式转换？

你想想看，PC 上的 Python 环境和 STM32 的嵌入式环境，完全是两个世界。TensorFlow 的 .h5 模型，STM32 读不懂；ONNX 格式，STM32Cube.AI 能消化一部分。但最终，我们要的是 .tflite 或者 Cube.AI 能直接导入的格式。

第一步：从 Keras 导出 .h5 模型

这个最简单，但坑也最多。我习惯在训练结束后，用 model.save() 保存完整模型，而不是只保存权重。为什么？因为部署时需要知道输入输出的 shape 和数据类型。

# 保存完整模型（包含网络结构 + 权重 + 训练配置）
model.save('my_model.h5')

# 或者只保存权重（不推荐用于部署）
# model.save_weights('my_weights.h5')

嗯，这里要注意一点：.h5 文件里包含了优化器状态和训练配置，这些在部署时完全没用。但没关系，Cube.AI 导入时会自动忽略它们。

第二步：导出 ONNX 格式

ONNX 是微软和 Facebook 搞的开放格式，好处是跨框架兼容。STM32Cube.AI 对 ONNX 的支持也越来越好。我个人建议，如果模型结构不太复杂，优先走 ONNX 路线。

转换需要用到 tf2onnx 工具：

# 安装工具
# pip install tf2onnx

import tf2onnx
import tensorflow as tf

# 加载 .h5 模型
model = tf.keras.models.load_model('my_model.h5')

# 指定输入签名（必须！）
spec = (tf.TensorSpec((None, 32, 32, 3), tf.float32, name="input"),)

# 转换
output_path = "my_model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, output_path=output_path)

print(f"ONNX 模型已保存至: {output_path}")

转换完成后，我建议用 onnxruntime 验证一下输出是否一致：

import onnxruntime as ort
import numpy as np

# 创建推理会话
sess = ort.InferenceSession("my_model.onnx")

# 构造随机输入
dummy_input = np.random.randn(1, 32, 32, 3).astype(np.float32)

# 推理
outputs = sess.run(None, {"input": dummy_input})
print("ONNX 推理成功，输出 shape:", outputs[0].shape)

第三步：模型量化与优化

这是最核心的一步。STM32 的算力和内存都有限，FP32 的模型直接放上去，大概率跑不动。量化就是把 FP32 的权重和激活值，压缩成 INT8 甚至 INT4。

量化的本质，说白了就是用精度换速度。我做过一个实验：同样的模型，FP32 推理一次要 120ms，INT8 量化后只要 18ms，快了将近 7 倍。代价是精度从 97.2% 降到了 96.5%，完全能接受。

量化方式对比

量化方式	精度损失	速度提升	适用场景
FP32（不量化）	无	1x	原型验证、资源充足的 MCU
动态范围量化	较小（0.5%~1%）	2~3x	快速部署、对精度敏感
全整数量化（INT8）	中等（1%~3%）	4~7x	大多数 STM32 部署场景
混合精度量化	可控	2~5x	部分层敏感、需要精细控制

实战：用 TFLite 做 INT8 量化

STM32Cube.AI 支持直接导入 TFLite 量化模型。我一般用 TensorFlow Lite Converter 来做：

import tensorflow as tf

# 加载模型
model = tf.keras.models.load_model('my_model.h5')

# 构造代表数据集（量化需要校准数据）
def representative_dataset():
    # 从训练集中取 100 张图片做校准
    for i in range(100):
        yield [train_images[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('my_model_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)

print("INT8 量化模型已保存")

知识结构总览

下面这张图，把整个模型准备与格式转换的流程串起来了。我建议你保存下来，每次做转换时对照着走一遍：

常见问题与避坑指南

• ONNX 转换失败：检查模型里有没有自定义层。如果有，需要先注册自定义算子，或者用 Keras 原生层替换。
• 量化后精度暴跌：八成是代表数据集没选好。确保校准数据覆盖了所有类别，数量不少于 100 张。
• Cube.AI 导入报错：检查输入输出 tensor 的名称。ONNX 模型里默认叫 "input" 和 "output"，但如果你改了名字，Cube.AI 里要对应设置。
• 模型太大放不进 Flash：考虑更激进的量化（INT8），或者剪枝后再量化。我有个项目从 2.3MB 压到 280KB，精度只掉了 1.2%。