资源受限 MCU 上的模型分片推理设计：Flash 与 RAM 之间按层换入换出的调度策略

资源受限 MCU 上的模型分片推理设计Flash 与 RAM 之间按层换入换出的调度策略一、128KB RAM 跑 1MB 模型的极限挑战当模型体量远超物理内存一块 STM32H743 MCU 拥有 1MB Flash 和约 512KB 的连续 RAM。一个量化后的 MobileNetV1 模型int8参数约占 4.2MB——即便使用外部 QSPI Flash 扩展存储空间片内 RAM 仍然只能容纳模型的一小部分。传统的做法是压缩再压缩通道剪枝、8bit 量化、甚至二值化。但精度下降到一定阈值后应用价值归零。分片推理提供了一条不同的路径不做模型结构修改而是将模型按层Layer切分为多个片段在推理时逐层从 Flash 加载权重和中间张量到 RAM计算完毕后将输出保留在 RAM 中作为下一层的输入。本文探讨的调度策略核心是利用层间数据依赖关系设计双缓冲换入换出机制将 Flash 的带宽转换为计算吞吐。二、按层分片调度的数据流模型前向传播中的生产者-消费者链卷积神经网络的前向传播天然形成一条线性流水线第 N 层的输出是第 N1 层的输入。这为分片调度提供了基础——只需要在任一时刻保留当前层权重 输入张量 输出张量三块数据。flowchart TD subgraph Flash[外部 QSPI Flash1~16MB] W1[Layer1 权重] W2[Layer2 权重] W3[Layer3 权重] WN[LayerN 权重] end subgraph RAM[片内 SRAM≤512KB] subgraph BUF_A[Buffer A输入张量] IN[input tensor] end subgraph BUF_B[Buffer B输出张量] OUT[output tensor] end subgraph BUF_W[Buffer W当前层权重] WEIGHT[current layer weights] end end Flash --|DMA 加载权重| WEIGHT Flash --|QSPI 读取| IN IN --|feed| COMPUTE[Conv/FC/Activation 计算] WEIGHT --|feed| COMPUTE COMPUTE --|store| OUT OUT -.-|下一轮时角色互换| IN W1 -- BUF_W W2 -- BUF_W W3 -- BUF_W WN -- BUF_W双缓冲策略的核心逻辑是推理过程中 BUF_A 和 BUF_B 在输入和输出角色之间交替。第 N 层计算完毕后BUF_A 可以被释放BUF_B 成为下一层的输入缓冲区BUF_A 转而用于容纳第 N1 层的输出。整个过程只需要两块活动张量缓冲区和一块权重缓冲区。三、调度器实现Flash 带宽感知的分块预取与 DMA 流水线以下实现针对 STM32H743 W25Q128QSPI Flash四线模式读取约 40MB/s的组合/* shard_inference.c —— MCU 分片推理调度器 * 目标平台STM32H743480MHz Cortex-M7 * 存储片内 SRAM (512KB) 外部 QSPI Flash (16MB) * 策略双缓冲 按层 DMA 预取 */ #include stdint.h #include string.h #include qspi_driver.h #include tensor_op.h /* ——— 可配置参数 ——— */ #define MAX_TENSOR_SIZE (128 * 1024) /* 单张量最大 128KB */ #define MAX_WEIGHT_BUF (256 * 1024) /* 权重缓冲区 256KB */ #define QSPI_SECTOR_SIZE 4096 /* Flash 擦除单位 */ /* ——— 双缓冲区结构 ——— */ typedef struct { float *buf_a; /* 输入/输出 缓冲区 A */ float *buf_b; /* 输入/输出 缓冲区 B */ float *buf_w; /* 权重缓冲区 */ size_t max_buf_size; /* 缓冲区容量字节 */ int current_input; /* 当前输入使用哪个 buf0A, 1B */ size_t input_size; /* 当前输入张量大小 */ } DoubleBuffer; /* ——— 模型分片描述表Flash 中的元数据 ——— */ typedef struct { uint32_t flash_offset; /* 权重在 QSPI Flash 中的偏移 */ uint32_t weight_size; /* 权重占用字节数 */ uint32_t input_shape[4]; /* 输入维度 [N,H,W,C] */ uint32_t output_shape[4];/* 输出维度 [N,H,W,C] */ uint8_t op_type; /* 算子类型0Conv, 1FC, 2Pool... */ uint8_t activation; /* 激活函数0ReLU, 1None... */ uint16_t kernel[2]; /* 卷积核 [h, w]仅 Conv 有效 */ } LayerSegment; /* * 初始化双缓冲区——从链接脚本中保留的 RAM 区域划分 * 这些区域在 .ld 文件中通过 __RAM_BUF_A__, __RAM_BUF_W__ 等符号定义 */ int dbuf_init(DoubleBuffer *db, float *buf_a, float *buf_b, float *buf_w, size_t buf_size) { if (!db || !buf_a || !buf_b || !buf_w || buf_size 0) return -1; // 参数校验失败 db-buf_a buf_a; db-buf_b buf_b; db-buf_w buf_w; db-max_buf_size buf_size; db-current_input 0; // 初始使用 buf_a 作为输入 db-input_size 0; return 0; } /* * 核心调度函数逐层执行分片推理 * * param db 已初始化的双缓冲区 * param segments 层描述数组位于 Flash 中 * param num_layers 层数 * param input 模型的总输入数据如 224x224x3 的图像 * return 0成功, -1缓冲区溢出, -2Flash 读取失败 */ int shard_inference_run(DoubleBuffer *db, const LayerSegment *segments, int num_layers, const float *input) { if (num_layers 0) return -1; /* 将模型输入放入缓冲区 A */ const LayerSegment *first segments[0]; size_t first_input_bytes first-input_shape[0] * first-input_shape[1] * first-input_shape[2] * first-input_shape[3] * sizeof(float); if (first_input_bytes db-max_buf_size) return -1; // 溢出 memcpy(db-buf_a, input, first_input_bytes); db-input_size first_input_bytes; db-current_input 0; // A 是输入 for (int layer 0; layer num_layers; layer) { const LayerSegment *seg segments[layer]; /* ——— 步骤1从 QSPI Flash 加载当前层权重 ——— */ if (seg-weight_size MAX_WEIGHT_BUF) return -1; int ret qspi_read_data(seg-flash_offset, (uint8_t *)db-buf_w, seg-weight_size); if (ret ! 0) { // Flash 读取超时或校验失败 return -2; } /* ——— 步骤2计算输出张量大小 ——— */ size_t out_size seg-output_shape[0] * seg-output_shape[1] * seg-output_shape[2] * seg-output_shape[3] * sizeof(float); if (out_size db-max_buf_size) return -1; // 输出超出缓冲区 /* ——— 步骤3确定输入/输出缓冲区指针 ——— */ float *cur_input db-current_input 0 ? db-buf_a : db-buf_b; float *cur_output db-current_input 0 ? db-buf_b : db-buf_a; /* ——— 步骤4执行本层计算 ——— */ switch (seg-op_type) { case OP_CONV: conv2d_naive(cur_input, // 输入 cur_output, // 输出 db-buf_w, // 权重 seg-input_shape, // in [N,H,W,C] seg-output_shape, // out [N,H,W,C] seg-kernel[0], seg-kernel[1]); // kernel h,w break; case OP_FC: fc_layer(cur_input, cur_output, db-buf_w, seg-input_shape, seg-output_shape); break; case OP_POOL: maxpool2d(cur_input, cur_output, seg-input_shape, seg-output_shape, seg-kernel[0], seg-kernel[1]); break; default: return -1; // 不支持的算子类型 } /* ——— 步骤5角色互换——当前输出变为下一层的输入 ——— */ db-current_input (db-current_input 0) ? 1 : 0; db-input_size out_size; } return 0; }四、边界分析QSPI 带宽瓶颈与算子粒度选择的工程权衡分片推理的性能天花板由 QSPI Flash 读取带宽决定。以 W25Q128 四线模式40MB/s为例算子类型权重规模Flash 加载时间计算时间ARM CMSIS-NN瓶颈Conv 3x3x32x6418KB0.46ms2.1ms计算FC 1024x5122MB51.2ms1.0msint8 矩阵乘Flash 读取Depthwise 3x3x640.6KB0.015ms0.8ms计算FC 层的场景下Flash 加载时间远超计算时间——此时使用 DMA 后台预取是关键优化手段。QSPI 控制器支持在计算当前层的同时通过 DMA 流水线将下一层的权重预加载到中间缓冲区。适用边界Flash 容量 模型权重总量的 1.2 倍留出校验与冗余空间计算/搬移比 0.8 的层带宽浪费可接受对于 FC 等大权重层应独立评估是否适合分片不可用于循环神经网络RNN/LSTM因为隐藏状态的时间依赖打破了层间独立性五、总结分片推理在 RAM 不到模型 1/10 的极端场景下提供了可行的执行方案。双缓冲机制利用前向传播的层间独立特性将以层为单位的权重加载和计算串行化将存储容量问题转化为带宽管理问题。关键优化方向是将 QSPI 的 DMA 预取与 CMSIS-NN 的计算并行执行以及根据 Flash 带宽特性合理拆分卷积层和全连接层的分片边界。在实际落地中应优先争取扩大片上 SRAM 以减少分片次数——每多一次 Flash 访问总延迟约增加 200us 的读命令开销。