基于STM32与PyTorch的端云协同AI案例：模型训练与轻量化部署-平芜编程栈

基于STM32与PyTorch的端云协同AI案例：模型训练与轻量化部署

1. 场景引入：当AI遇见嵌入式设备

想象一下这样的场景：一个智能家居设备需要实时识别用户说出的控制指令，比如"开灯"、"调高温度"或"播放音乐"。传统方案要么依赖云端处理（带来延迟和隐私问题），要么受限于嵌入式设备的算力难以实现。这正是端云协同AI大显身手的地方。

在这个案例中，我们将展示如何用PyTorch训练一个轻量级关键词识别模型，经过优化后部署到STM32微控制器上。云端负责复杂的模型训练和迭代，终端设备则实现低功耗、实时的本地推理。这种架构既保护了用户隐私，又确保了响应速度，是边缘计算的典型应用。

2. 解决方案设计

2.1 整体架构

我们的方案采用经典的端云协同架构：

云端：使用PyTorch 2.8镜像训练和优化模型
- 数据收集与标注
- 模型训练与验证
- 模型轻量化（量化+剪枝）
端侧：STM32微控制器部署
- 模型转换与部署
- 实时音频采集与预处理
- 本地推理与响应

2.2 为什么选择STM32

STM32系列微控制器在嵌入式领域应用广泛，具有以下优势：

低功耗设计，适合长时间运行的IoT设备
丰富的外设接口，便于连接各种传感器
充足的存储空间（本例使用STM32H7系列，具有1MB Flash和512KB RAM）
成熟的开发工具链支持

3. 云端模型开发

3.1 数据准备

我们使用开源语音命令数据集，包含30个关键词的65000条1秒语音样本。数据预处理流程如下：

import torchaudio from torchaudio.transforms import MelSpectrogram def preprocess_audio(waveform): # 转换为梅尔频谱图 transform = MelSpectrogram( sample_rate=16000, n_mels=40, n_fft=1024, hop_length=256 ) return transform(waveform)

3.2 模型设计与训练

我们采用轻量化的CNN架构，适合嵌入式部署：

import torch.nn as nn class KeywordCNN(nn.Module): def __init__(self, num_classes=30): super().__init__() self.features = nn.Sequential( nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(2) ) self.classifier = nn.Sequential( nn.Linear(32 * 10 * 10, 64), nn.ReLU(), nn.Linear(64, num_classes) ) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) x = self.classifier(x) return x

训练过程使用PyTorch 2.8的混合精度训练加速：

model = KeywordCNN().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) scaler = torch.cuda.amp.GradScaler() for epoch in range(50): for inputs, labels in train_loader: inputs, labels = inputs.cuda(), labels.cuda() with torch.cuda.amp.autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

3.3 模型轻量化

为适应STM32的有限资源，我们对模型进行优化：

量化：将FP32模型转换为INT8，减少75%的存储和计算需求

quantized_model = torch.quantization.quantize_dynamic( model, {nn.Linear}, dtype=torch.qint8 )

剪枝：移除不重要的连接，减少20%的参数

parameters_to_prune = [ (module, 'weight') for module in filter( lambda m: isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear), model.modules() ) ] prune.global_unstructured( parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.2 )

优化后模型大小从1.2MB降至280KB，准确率仅下降2.3%，完全满足嵌入式部署要求。

4. 端侧部署

4.1 模型转换

使用STM32Cube.AI工具将PyTorch模型转换为STM32可执行格式：

stm32ai convert -m quantized_model.onnx -v 0 --allocate-inputs --output-dir ./stm32_model

4.2 STM32工程集成

在STM32CubeIDE中创建工程，关键配置如下：

启用I2S接口用于音频采集
配置DMA实现高效数据传输
集成STM32Cube.AI运行时库
分配模型所需的Flash和RAM空间

推理代码示例（基于STM32Cube HAL库）：

// 初始化AI运行时 ai_handle_t network = AI_HANDLE_NULL; ai_buffer_t input_buffer, output_buffer; ai_error err = ai_mnetwork_create(&network, AI_MNETWORK_DATA_CONFIG); if (err != AI_ERROR_NONE) { Error_Handler(); } // 音频采集与预处理 I2S_Receive(&hi2s2, (uint16_t*)pcm_buffer, AUDIO_BUFFER_SIZE); // 执行推理 input_buffer.data = AI_PTR(preprocessed_audio); ai_mnetwork_run(network, &input_buffer, &output_buffer); // 解析结果 uint8_t predicted_class = argmax((float*)output_buffer.data);