PyTorch 模型部署到 STM32N6

这篇文章解决什么问题

你已经有 PyTorch 模型，现在要把它放到 STM32N6 上跑。
本文给你一条可落地路径：

• PC 端模型怎么准备；
• 为什么要导出 ONNX；
• ST Edge AI Core 在链路里做什么；
• CubeMX + VS Code 工程怎么接；
• 板端结果和 PC 不一致时从哪里查。

默认工具链：

• STM32CubeMX
• STM32CubeIDE for VS Code
• STM32CubeProgrammer
• ST Edge AI Core

---

先理解部署链路（背后原理）

你在 PC 上的 PyTorch 模型，不能直接原样放到 MCU 里跑，原因是：

• 训练框架依赖重（Python 运行时、算子实现复杂）；
• MCU 资源有限（RAM/Flash/带宽）；
• 嵌入式推理需要静态、可裁剪、可编译的图结构。

所以常见链路是：

1. PyTorch 导出为 ONNX（静态计算图交换格式）
2. ST Edge AI Core 分析并转换成目标平台可执行形式
3. 生成/集成 C 工程代码，在 STM32N6 上运行

核心思想：把“训练图”变成“部署图”。

---

第一步：在 PC 端把模型“部署化”

操作

先固定模型行为和输入尺寸：

import torch

model = MyModel()
model.load_state_dict(torch.load("best.pt", map_location="cpu"))
model.eval()

# 固定输入 shape（示例）
dummy = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=13,
    do_constant_folding=True,
)
print("Exported model.onnx")

原理

• model.eval()：关闭 Dropout、固定 BN 统计，保证推理稳定；
• 固定输入 shape：让后端更容易做内存规划和算子映射；
• ONNX：作为框架间的中间表示，便于后续工具链接入。

常见坑

• 动态输入 shape 导致后端部署困难；
• 使用冷门算子导致转换失败；
• 训练时预处理和部署时预处理不一致。

---

第二步：用 ST Edge AI Core 做分析和转换

操作

先确认工具可用：

stedgeai --version
stedgeai --help

然后按你本机版本的命令执行分析和生成流程（不同版本参数名可能不同，优先 --help）：

• 分析：看算子支持、RAM/Flash 估算、性能估算；
• 生成：导出可集成进 STM32 工程的产物。

原理

ST Edge AI Core 的作用不是“训练模型”，而是做部署编译：

• 把通用模型图转换成目标平台支持的执行图；
• 做算子替换、图优化、内存规划；
• 产出报告（性能、内存、精度风险）和工程集成文件。

量化说明（INT8）

量化本质是用更低位宽表示参数和激活，换速度/内存，可能损失精度。
所以要有校准集，并做量化前后对比，不要直接盲上板。

---

第三步：在 STM32N6 工程里接入推理

操作

1. 用 CubeMX 建立 STM32N6 工程，完成时钟和基础外设；
2. 在 AI 相关配置里导入转换产物；
3. 生成工程，使用 STM32CubeIDE for VS Code 打开；
4. 在代码中串起预处理 -> 推理 -> 后处理。

最小主循环结构：

while (1)
{
    // 1) 采集或准备输入
    // 2) preprocess: 归一化、resize、通道排列

    // 3) 调用推理入口（函数名以生成代码为准）
    // ai_run(ai_input, ai_output);

    // 4) postprocess: 阈值/NMS/类别解码

    // 5) 串口打印结果和耗时
}

原理

把流程拆成四段的原因：

• 预处理错了，模型再好也输出垃圾；
• 推理层稳定后，问题可快速定位到前后处理；
• 后处理是检测类任务误差高发区（阈值、NMS、坐标变换）。

---

第四步：VS Code 编译下载与“对拍验证”

操作

• 编译工程；
• ST-Link 下载到板子；
• 串口输出推理结果与单次耗时；
• 用同一组输入样本，让 PC 与板端对拍。

原理

“编译通过”不等于部署成功。部署成功的判据是：

• 功能：输出结果正确；
• 一致性：与 PC 参考接近；
• 性能：延时满足需求；
• 稳定性：连续运行不崩。

---

第五步：出现问题时怎么排（操作 + 原理）

1) 板子跑崩或 HardFault

先查：

• 输入输出 buffer 大小是否和模型要求一致；
• 大数组放置是否合理（避免栈爆）；
• 是否有越界访问。

原因：嵌入式端内存紧，越界和栈溢出比 PC 更容易直接崩溃。

2) 输出和 PC 明显不一致

先查预处理：

• 归一化范围（01 还是 -11）；
• 均值方差；
• RGB/BGR 通道顺序；
• resize 方式和插值方法。

原因：模型学到的是“训练输入分布”，不是你的想当然输入。

3) 延时太高

优先动作：

• 降输入分辨率；
• 换轻量模型；
• 做量化；
• 减少后处理复杂度。

原因：MCU/NPU 算力和内存带宽有限，检测类后处理常是隐藏瓶颈。

4) 串口没结果

先做最小化：

• 先打印固定字符串确认串口通；
• 再打印推理前后时间戳；
• 最后再看模型逻辑。

原因：很多“模型问题”本质是 IO 或流程没走到。

---

你们这套 STM32N6 项目的实用建议

如果是首次落地，建议顺序：

1. 先用小分类模型跑通全链路；
2. 再上检测模型（如 YOLO）；
3. 再做精度和延时优化。

原因很直接：

• 分类模型链路短，便于验证工具链和工程结构；
• 检测模型会引入更多后处理和内存压力，先跑通再加复杂度更稳。

---

一句话结尾

STM32N6 部署最容易失败的点不是“模型不够强”，而是“输入定义和工程链路不一致”。
先把导出、转换、集成、对拍这四步做成固定流水线，你后面换模型会轻松很多。