PyTorch 将 ExecuTorch 扩展至 Arm Cortex-M:实现边缘端模型原生部署
背景:打破 MCU 上的 PyTorch 部署壁垒
微控制器(MCU)作为物联网、可穿戴设备和工业传感器的核心,正面临日益增长的机器学习推理需求。然而,由于资源极度受限(通常小于 1 MB SRAM 和数兆字节 Flash),在 Cortex-M 架构上部署模型曾是一个难题。
过去,开发者若使用 PyTorch 进行训练,必须将模型转换为其他框架格式(如 TensorFlow Lite 或 ONNX),并维护独立的工具链。这不仅引入了跨框架的数值不匹配风险,还增加了迭代周期中的摩擦成本。对于在顶级 ML 会议上发布超过 90% 研究代码的 PyTorch 团队而言,这种转换步骤是巨大的障碍。
核心突破:Cortex-M 后端与原生支持
PyTorch 推出的 ExecuTorch Cortex-M 后端 彻底改变了这一局面。它允许 PyTorch 模型直接在 Arm Cortex-M CPU 上运行,并无缝利用 CMSIS-NN(Arm 开源优化的神经网络内核库)进行加速。
1. 全能力层级覆盖
新后端支持 Cortex-M 处理器的所有能力层级,通过单一工作流覆盖不同硬件特性:
- Scalar 层级:适用于 M0, M0+, M3 等超低成本芯片。
- DSP 层级:利用 SIMD 乘累加指令,适用于 M4, M7 等中端 MCU。
- MVE (Helium) 层级:利用 128 位向量操作和基于拍(beat)的执行模型,在 M55, M85 等高性能芯片上提供最大性能增益。
2. 极致轻量与编译优化
ExecuTorch 的核心运行时仅占 20 KiB Flash,且采用按需编译(Selective Build)策略,仅链接模型所需的算子,代码规模随模型而非框架线性增长。
在编译阶段,ExecuTorch 执行了一系列针对 Cortex-M 的优化:
- 图级计算消除:移除冗余的数据类型转换、量化/反量化对以及内存格式转置。
- 量化参数折叠:将缩放和零点算术预计算为整数乘法和移位,确保量化核推理完全在定点数域内完成。
- 权重布局转置:在导出时将 PyTorch 的 OIHW 布局转换为 CMSIS-NN 期望的 OHWI 布局,避免运行时开销。
- 静态内存规划:支持跨异构内存(如 TCM 和 SRAM)分配,无堆分配器,无碎片,无运行时内存管理。
实测数据:性能与效率双优
团队在 Alif Ensemble E8(搭载 Cortex-M55 核心)上对 MLPerf Tiny 套件中的四个模型进行了基准测试,结果令人印象深刻:
- 运行时间:所有模型运行时间均 ≤ 21.5 ms。
- 内存占用:单核 Cortex-M55 内存占用 < 90 KiB。
- 异构 SoC 支持:对于包含 Ethos-U NPU 的异构 SoC,该后端提供了 CPU 执行路径与 NPU 委托的无缝切换。
价值总结
ExecuTorch 的 Cortex-M 后端不仅填补了 PyTorch 在边缘端的空白,更通过“一次训练,多处部署”的理念,大幅降低了边缘 AI 开发的门槛。开发者现在可以在同一工具链中完成从云端训练到微控制器推理的全流程,无需担心格式转换带来的数值误差或工具链割裂。
“ExecuTorch 消除了这一转换步骤,使 PyTorch 模型能够直接在 Arm Cortex-M CPU 上运行,并利用 CMSIS-NN 加速。它支持从 Scalar 到 Helium 的所有能力层级,并在单一工作流中完成。” —— PyTorch 官方团队