2. 手写数字识别(MNIST)
任务目标
本案例将完成一个端到端的MNIST手写数字识别部署流程:使用Keras训练一个轻量级CNN模型,通过STM32Cube.AI将其转换为STM32可推理的C代码,最终在开发板上对28×28灰度手写数字图片进行分类,并将识别结果(0-9)通过串口或LCD显示。
环境与依赖
- PC端:Python 3.8+,TensorFlow 2.4+,STM32CubeMX 6.4+(含Cube.AI扩展包),STM32CubeIDE 1.8+
- 硬件:STM32F746G-DISCO开发板(或任意支持Cube.AI的Cortex-M7/M4板卡)
- 模型输入:28×28×1 float32(归一化至[0,1])
- 模型输出:10类softmax概率分布
模型训练与转换
首先在PC端训练一个参数量约6K的CNN模型(两个卷积层+两个全连接层),量化感知训练后导出为Keras H5文件。随后在STM32CubeMX中导入该H5文件,Cube.AI将自动生成network.c、network.h及network_data.c,其中包含量化后的权重和推理函数ai_run。
关键完整代码
以下为STM32端核心推理与显示逻辑(省略Cube.AI生成的网络层代码,仅展示调用入口):
/* main.c 关键片段 */
#include "network.h"
#include "ai_platform.h"
#include "stdio.h"
/* Cube.AI生成的网络实例 */
AI_NETWORK_DECLARE(network);
static ai_handle network_handle = AI_HANDLE_NULL;
static ai_buffer input_buffers[1];
static ai_buffer output_buffers[1];
/* 28x28 灰度图像缓冲区(由摄像头或串口填充) */
static float img_data[28*28];
/* 初始化Cube.AI网络 */
void ai_init(void) {
ai_error err;
const ai_network_params params = {
AI_NETWORK_DATA_WEIGHTS(ai_network_data_weights_get()),
AI_NETWORK_DATA_ACTIVATIONS(ai_network_data_activations_get())
};
err = ai_network_create(&network_handle, &network);
if (err.type != AI_ERROR_NONE) {
printf("Network create failed: %d\n", err.code);
while(1);
}
if (ai_network_init(network_handle, ¶ms) != AI_ERROR_NONE) {
printf("Network init failed\n");
while(1);
}
input_buffers[0] = ai_network_inputs_get(network_handle, NULL);
output_buffers[0] = ai_network_outputs_get(network_handle, NULL);
}
/* 执行推理并返回预测数字 */
uint8_t ai_predict(float* img) {
ai_i8 quantized_input[28*28];
ai_i8 quantized_output[10];
float output[10];
/* 输入量化:float [0,1] → int8 [-128,127] */
ai_float_to_i8(img, quantized_input, 28*28, 1.0f/128.0f, -128.0f);
input_buffers[0].data = quantized_input;
output_buffers[0].data = quantized_output;
if (ai_run(network_handle, input_buffers, output_buffers) != AI_ERROR_NONE) {
return 0xFF; /* 推理失败 */
}
/* 输出反量化 */
ai_i8_to_float(quantized_output, output, 10, 1.0f/128.0f, -128.0f);
/* 取最大概率索引 */
uint8_t max_idx = 0;
for (int i = 1; i < 10; i++) {
if (output[i] > output[max_idx]) max_idx = i;
}
return max_idx;
}
/* 主循环示例:从串口接收28*28字节,推理并回传结果 */
int main(void) {
HAL_Init();
SystemClock_Config();
MX_GPIO_Init();
MX_USART1_UART_Init();
ai_init();
printf("MNIST Inference Ready\r\n");
while (1) {
/* 接收28*28=784字节原始像素(0-255) */
uint8_t raw[784];
HAL_UART_Receive(&huart1, raw, 784, 5000);
/* 归一化至[0,1] */
for (int i = 0; i < 784; i++) {
img_data[i] = raw[i] / 255.0f;
}
uint8_t digit = ai_predict(img_data);
printf("Predicted: %d\r\n", digit);
}
}
运行与验证
- PC端训练:运行训练脚本生成
mnist_cnn.h5,验证测试集准确率≥98% - CubeMX配置:选择STM32F746G-DISCO,在Middleware中启用Cube.AI,导入H5文件,生成代码
- 编译烧录:将上述代码合并至生成的工程,编译后烧录至开发板
- 验证方法:通过串口助手发送784字节的MNIST测试图片(可从测试集提取),观察返回的数字是否与标签一致
通过标准:连续测试10张不同数字的图片,识别正确≥9张即判定部署成功。若使用LCD显示,需额外添加显示驱动,本案例以串口输出为验证路径。
常见问题
- 内存不足:Cube.AI生成的激活缓冲区过大时,在CubeMX中减小输入批次或使用更深量化(int8)
- 精度下降:量化后准确率下降超过1%时,在训练阶段加入伪量化节点(QAT)
- 串口数据格式:发送的784字节需为原始像素值(0-255),且按行优先排列