当前位置:首页 > 嵌入式 > 嵌入式分享
[导读]TinyML的开发流程存在一个天然的断裂带:数据科学家习惯使用PyTorch等框架在云端GPU上训练模型,而嵌入式工程师则需要在Keil、Arduino或ESP-IDF环境中编写C++代码。这种技术栈的割裂导致模型从训练到部署往往需要数周的手工重写和调试。跨平台迁移的核心理念是建立一条自动化的转换流水线,让PyTorch训练的模型能够无损地运行在STM32和ESP32这类资源受限的微控制器上。本文将系统阐述从模型导出、格式转换到嵌入式集成的完整流程,并提供可复现的工程实践方案。

TinyML的开发流程存在一个天然的断裂带:数据科学家习惯使用PyTorch等框架在云端GPU上训练模型,而嵌入式工程师则需要在Keil、Arduino或ESP-IDF环境中编写C++代码。这种技术栈的割裂导致模型从训练到部署往往需要数周的手工重写和调试。跨平台迁移的核心理念是建立一条自动化的转换流水线,让PyTorch训练的模型能够无损地运行在STM32和ESP32这类资源受限的微控制器上。本文将系统阐述从模型导出、格式转换到嵌入式集成的完整流程,并提供可复现的工程实践方案。

从PyTorch到ONNX:模型导出的第一步

PyTorch模型无法直接被TinyML推理引擎识别,需要首先导出为开放的中间表示格式。ONNX是当前最成熟的跨框架模型交换标准,它通过计算图的形式保存模型结构和参数,与原始框架解耦。导出的关键在于处理动态控制流和输入张量的维度信息。

在PyTorch中,模型导出需要执行一次虚拟推理来记录计算图。以下代码演示了将训练好的图像分类模型导出为ONNX格式的标准流程:

import torch

import torch.onnx

# 加载预训练的PyTorch模型并设置为推理模式

model = torch.load('my_model.pth')

model.eval()

# 创建示例输入(批次大小1,3通道,224x224图像)

dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX模型

torch.onnx.export(

model,

dummy_input,

'model.onnx',

input_names=['input'],

output_names=['output'],

dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}},

opset_version=11

)

导出后的ONNX文件包含了完整的计算图定义,可以使用Netron等可视化工具查看模型结构,验证各层连接是否正确。这一步骤的产物将成为后续跨平台转换的通用输入。

## ONNX到TensorFlow Lite的格式转换

TinyML生态的通用语言是TensorFlow Lite格式。无论是STM32Cube.AI还是ESP-TFLite-Micro,都原生支持.tflite文件的解析和执行。因此,ONNX到TFLite的转换是整个迁移流程的核心环节。

这一转换通常需要借助ONNX-TensorFlow转换工具完成两步操作:先将ONNX转换为TensorFlow的冻结图格式,再利用TFLite转换器生成轻量级模型。转换过程中最关键的是量化配置——将32位浮点权重转换为8位整数,这是模型能够在MCU上运行的前提条件。

import onnx

from onnx_tf.backend import prepare

import tensorflow as tf

# 加载ONNX模型

onnx_model = onnx.load('model.onnx')

tf_rep = prepare(onnx_model)

# 导出为TensorFlow冻结图

tf_rep.export_graph('model_frozen.pb')

# 配置TFLite转换器并启用量化

converter = tf.lite.TFLiteConverter.from_frozen_graph(

'model_frozen.pb', ['input'], ['output']

)

converter.optimizations = [tf.lite.Optimize.DEFAULT]

converter.representative_dataset = representative_dataset # 校准数据集

converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

# 执行转换

tflite_model = converter.convert()

with open('model_quantized.tflite', 'wb') as f:

f.write(tflite_model)

量化后的模型体积可压缩至原来的四分之一,同时整数运算在Cortex-M内核上的执行效率远高于浮点运算。

## STM32平台的集成:STM32Cube.AI工作流

STM32生态系统提供了STM32Cube.AI作为官方模型转换工具。该工具作为STM32CubeMX的扩展插件,能够将.tflite或ONNX模型自动转换为针对目标STM32芯片优化的C代码库。转换过程不仅生成模型权重数组,还封装了输入输出缓冲区的内存管理和推理调度逻辑。

部署到STM32的标准流程如下:首先在STM32CubeMX中创建工程并选择目标芯片(如STM32H7系列),在Middleware设置中启用X-CUBE-AI组件并导入.tflite模型文件。工具会自动分析模型结构,报告RAM和Flash的预估占用。生成代码后,在用户工程中调用AI库提供的接口函数即可执行推理:

#include "ai_platform.h"

#include "network.h" // STM32Cube.AI生成的模型头文件

// 声明输入输出缓冲区

AI_ALIGNED(4) static ai_u8 activations[AI_NETWORK_DATA_ACTIVATIONS_SIZE];

static ai_handle network;

// 初始化模型

ai_network_create(&network, AI_NETWORK_DATA_CONFIG);

ai_network_init(network, activations, sizeof(activations));

// 准备输入数据(如传感器采集的128点FFT结果)

ai_buffer* input = ai_network_inputs_get(network, NULL);

memcpy(input->data, sensor_data, input->size * sizeof(ai_float));

// 执行推理

ai_network_run(network, &input, &output);

// 读取输出结果

ai_buffer* output = ai_network_outputs_get(network, NULL);

float confidence = output->data[0];

STM32Cube.AI的优势在于与CubeMX生态的深度集成,开发者可以使用相同的工具链完成外设配置和AI模型部署,学习成本较低。

## ESP32平台的集成:ESP-TFLite-Micro组件

ESP32平台的TinyML部署由乐鑫官方提供的esp-tflite-micro组件支撑。该组件是TensorFlow Lite Micro框架的ESP-IDF移植版本,并集成了ESP-NN加速库,针对ESP32系列芯片的SIMD指令集进行了深度优化。

在ESP-IDF工程中安装该组件只需一条命令:

idf.py add-dependency "espressif/esp-tflite-micro^1.3.1"

组件的使用接口与标准TFLite Micro保持一致,核心代码包含模型加载、张量分配和推理执行三个步骤。ESP-NN加速库的效果极为显著:在ESP32-S3上执行人员检测模型时,启用ESP-NN后推理时间从2300毫秒缩短至54毫秒,性能提升超过40倍。

#include "tensorflow/lite/micro/all_ops_resolver.h"

#include "tensorflow/lite/micro/micro_interpreter.h"

#include "model.h" // xxd转换生成的模型数组

// 定义内存池

constexpr int kTensorArenaSize = 50 * 1024;

static uint8_t tensor_arena[kTensorArenaSize];

// 加载模型

const tflite::Model* model = tflite::GetModel(g_model);

static tflite::AllOpsResolver resolver;

static tflite::MicroInterpreter interpreter(model, resolver, tensor_arena,

kTensorArenaSize);

// 分配张量并获取输入指针

interpreter.AllocateTensors();

float* input = interpreter.input(0)->data.f;

// 执行推理

interpreter.Invoke();

float* output = interpreter.output(0)->data.f;

对于Arduino开发环境的用户,Seeed Studio提供了将TFLite Micro集成到Arduino IDE的完整方案,包含数据采集、模型训练和hex文件烧录的全流程指引。

## 实际案例:关键词识别模型的跨平台迁移

以语音唤醒应用为例,完整的迁移流程可以验证上述方法的可行性。首先在PyTorch中训练一个基于深度可分离卷积的关键词识别模型,输入为40维MFCC特征,输出为4个类别。导出为ONNX后,经TFLite转换生成INT8量化的.tflite文件,体积约35KB。

在STM32F407平台,STM32Cube.AI生成的代码占用约42KB RAM和48KB Flash,单次推理耗时约22毫秒。在ESP32-S3平台,使用esp-tflite-micro组件并启用ESP-NN加速,推理延迟约为18毫秒,RAM占用约38KB。两个平台的推理精度与原始PyTorch模型相比,下降幅度均在0.5%以内,完全满足实际应用需求。

结语

从PyTorch到STM32/ESP32的模型迁移已经形成了一条清晰的技术路径:PyTorch导出ONNX,ONNX转换TFLite,TFLite通过厂商工具链生成嵌入式C代码。这条路径上的每个环节都有成熟的开源工具支撑,开发者无需深入每个框架的内部实现即可完成端到端部署。关键的成功要素在于量化配置的正确设置和厂商加速库的充分使用。随着TinyML生态的持续成熟,跨平台迁移的门槛将进一步降低,使边缘AI应用的开发更加敏捷和高效。

本站声明: 本文章由作者或相关机构授权发布,目的在于传递更多信息,并不代表本站赞同其观点,本站亦不保证或承诺内容真实性等。需要转载请联系该专栏作者,如若文章内容侵犯您的权益,请及时联系本站删除( 邮箱:macysun@21ic.com )。
换一批
延伸阅读
关闭