📚 课程目录第4章 MNIST手写数字识别

4. MNIST手写数字识别

任务目标

在STM32上部署TensorFlow Lite for Microcontrollers（TFLM）MNIST手写数字识别模型，通过串口接收28×28像素灰度图像数据，输出0~9识别结果。

环境与依赖

硬件：STM32F746G-DISCO开发板（或任意Cortex-M7/M4，≥320KB RAM）
软件：STM32CubeIDE 1.13+，X-CUBE-AI 7.3.0（或手动集成TFLM库）
模型文件：mnist_model.tflite（4KB，量化后int8权重）
输入格式：28×28=784字节，灰度值0~255，行优先排列

实现方案

使用X-CUBE-AI将预训练MNIST模型转换为STM32可调用的C函数。模型输入为784个int8值（像素值减128后映射到[-128,127]），输出为10个类别得分。通过UART2接收784字节图像数据，调用AI推理函数，将得分最高的类别通过UART发送回上位机。

关键完整代码

/* main.c 核心逻辑 */
#include "ai_platform.h"
#include "network.h"
#include "network_data.h"
#include <stdio.h>
#include <string.h>

#define IMG_SIZE 784
#define UART_BUF_SIZE IMG_SIZE + 2

ai_handle network;
ai_network_report report;
ai_input ai_in[1];
ai_output ai_out[1];
int8_t input_data[IMG_SIZE];
int8_t output_data[10];
uint8_t uart_buf[UART_BUF_SIZE];
uint8_t uart_idx = 0;

/* 模型初始化 */
void ai_init(void) {
    ai_error err;
    err = ai_network_create(&network, AI_NETWORK_DATA_CONFIG);
    if (err.type != AI_ERROR_NONE) {
        Error_Handler();
    }
    ai_network_get_info(network, &report);
    
    ai_in[0].data = AI_HANDLE_PTR(input_data);
    ai_in[0].fmt = AI_NETWORK_DATA_FORMAT;
    ai_in[0].size = AI_NETWORK_IN_1_SIZE;
    ai_in[0].shape[0] = 1;
    
    ai_out[0].data = AI_HANDLE_PTR(output_data);
    ai_out[0].fmt = AI_NETWORK_DATA_FORMAT;
    ai_out[0].size = AI_NETWORK_OUT_1_SIZE;
    ai_out[0].shape[0] = 1;
}

/* 推理函数 */
int run_inference(void) {
    ai_i32 n_batch;
    ai_error err;
    
    /* 像素值偏移：uint8(0~255) → int8(-128~127) */
    for (int i = 0; i < IMG_SIZE; i++) {
        input_data[i] = (int8_t)(uart_buf[i] - 128);
    }
    
    n_batch = ai_network_run(network, ai_in, ai_out);
    if (n_batch != 1) {
        return -1;
    }
    
    /* 找最大得分索引 */
    int max_idx = 0;
    for (int i = 1; i < 10; i++) {
        if (output_data[i] > output_data[max_idx]) {
            max_idx = i;
        }
    }
    return max_idx;
}

/* UART接收回调（中断方式） */
void HAL_UART_RxCpltCallback(UART_HandleTypeDef *huart) {
    if (huart == &huart2) {
        uart_idx++;
        if (uart_idx < UART_BUF_SIZE) {
            HAL_UART_Receive_IT(&huart2, &uart_buf[uart_idx], 1);
        } else {
            /* 收到完整784字节+2字节帧头(0xAA55) */
            if (uart_buf[0] == 0xAA && uart_buf[1] == 0x55) {
                int result = run_inference();
                char msg[16];
                sprintf(msg, "Digit: %d\r\n", result);
                HAL_UART_Transmit(&huart2, (uint8_t*)msg, strlen(msg), 100);
            }
            uart_idx = 0;
            HAL_UART_Receive_IT(&huart2, &uart_buf[0], 1);
        }
    }
}

int main(void) {
    HAL_Init();
    SystemClock_Config();
    MX_GPIO_Init();
    MX_USART2_UART_Init();
    
    ai_init();
    HAL_UART_Receive_IT(&huart2, &uart_buf[0], 1);
    
    while (1) {
        /* 主循环空闲，推理在中断中完成 */
    }
}

运行与验证

在STM32CubeIDE中创建工程，启用X-CUBE-AI，导入mnist_model.tflite，生成代码后替换上述main.c核心逻辑
编译烧录，打开串口助手（115200-8-N-1）
发送数据帧：AA 55 + 784字节图像数据（行优先，灰度0~255）
预期返回格式：Digit: X\r\n，其中X为0~9

验证通过条件：使用MNIST测试集任选10张图像，逐张发送，识别正确率≥90%（即至少9张正确）。若使用STM32F746G-DISCO板载摄像头，可将摄像头采集的28×28灰度图按相同帧格式发送。

常见问题

RAM不足：确认模型已量化（int8），若仍不足，使用AI_NETWORK_DATA_ACTIVATIONS_SIZE宏检查激活内存，必要时减小输入批处理
识别错误：检查像素偏移是否正确（uint8转int8需减128），以及图像是否反色（MNIST白字黑底）
UART丢帧：确保上位机发送间隔>10ms，或使用DMA接收