3. 模型准备与ONNX/TFLite格式转换

模型训练好了，但离跑在STM32上还差关键一步——格式转换。说白了，Keras和PyTorch训练出来的模型，STM32Cube.AI不认识。你得先把它翻译成ONNX或TFLite这种中间格式。

我个人习惯把这一步叫做“模型护照办理”。没有护照，模型出不了训练环境，进不了嵌入式世界。今天我们就来搞定这件事。

3.1 为什么需要中间格式？

你想想看，PyTorch的模型结构是动态图，Keras是静态图，而STM32Cube.AI只认它自己的格式。直接硬转肯定不行。中间格式就像通用语言——ONNX是微软牵头搞的，TFLite是Google的移动端标准。STM32Cube.AI对这两种都支持得不错。

我在项目中遇到过一个问题：团队里有人用PyTorch，有人用Keras。最后统一转成ONNX，谁都不用迁就谁。这就是中间格式的价值。

3.2 PyTorch模型导出为ONNX

PyTorch转ONNX其实很直接。核心就一个函数——torch.onnx.export()。但有几个坑你得注意。

import torch
import torch.onnx

# 假设你有一个训练好的PyTorch模型
model = YourTrainedModel()
model.eval()  # 这一步不能省

# 构造一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出为ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 虚拟输入
    "model.onnx",        # 输出文件名
    export_params=True,  # 保存训练好的参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)
print("ONNX导出完成")

这里有个细节——opset_version。我建议用11或12。版本太低有些算子不支持，版本太高STM32Cube.AI可能还没跟上。嗯，11是个稳妥的选择。

3.3 Keras模型导出为TFLite

Keras转TFLite更简单。TensorFlow官方提供了TFLiteConverter，一行代码就能搞定。

import tensorflow as tf

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

# 转换为TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 保存为.tflite文件
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)
print("TFLite导出完成")

3.4 验证导出模型的正确性

导出完了不等于万事大吉。你得验证一下——导出的ONNX/TFLite和原始模型的输出是不是一致的。这一步很多人跳过，结果部署上去才发现推理结果不对，回头查问题特别痛苦。

我曾经有一次，模型导出后精度掉了5%，查了两天才发现是opset_version太低，某些算子被降级了。从那以后，我每次导出必做验证。

3.4.1 验证ONNX模型

import onnx
import onnxruntime
import numpy as np

# 1. 检查ONNX模型结构
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过")

# 2. 用ONNX Runtime跑推理
ort_session = onnxruntime.InferenceSession("model.onnx")

# 构造输入（和导出时一致）
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outputs = ort_session.run(None, ort_inputs)

# 3. 和原始PyTorch模型对比
with torch.no_grad():
    torch_output = model(torch.from_numpy(input_data))

# 比较输出差异
diff = np.abs(ort_outputs[0] - torch_output.numpy()).max()
print(f"最大差异: {diff:.6f}")

if diff < 1e-4:
    print("✅ 验证通过！ONNX和PyTorch输出一致")
else:
    print("❌ 差异过大，请检查导出过程")

3.4.2 验证TFLite模型

import tensorflow as tf
import numpy as np

# 加载TFLite模型
interpreter = tf.lite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()

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

# 构造输入
input_data = np.random.randn(1, 224, 224, 3).astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
tflite_output = interpreter.get_tensor(output_details[0]['index'])

# 和原始Keras模型对比
keras_output = model.predict(input_data, verbose=0)
diff = np.abs(tflite_output - keras_output).max()
print(f"最大差异: {diff:.6f}")

if diff < 1e-4:
    print("✅ TFLite验证通过")
else:
    print("❌ 请检查量化或导出参数")

3.5 常见问题与避坑指南

• 算子不支持：ONNX导出时如果遇到不支持的算子，可以试试降低opset_version，或者用torch.onnx.export的custom_opsets参数注册自定义算子。
• 动态输入尺寸：STM32上的模型输入尺寸通常是固定的。导出时建议固定batch_size=1，不要用动态轴。
• 量化后精度下降：如果TFLite量化后精度掉得厉害，试试representative_dataset——用一小部分真实数据校准量化参数。

3.6 本章知识结构

下面这张图帮你理清整个流程。从训练框架到中间格式，再到验证环节，每一步都不能少。