4. 基于CIFAR-10的简单CNN部署

任务目标

在STM32F407单片机上部署一个针对CIFAR-10数据集的简单CNN分类器，实现10类图像的实时推理。本案例聚焦于将训练好的CNN模型权重转换为C数组，并通过CMSIS-NN库完成前向推理，最终在串口输出分类结果。

环境依赖

• STM32F407开发板（192KB RAM，1MB Flash）
• STM32CubeIDE或Keil MDK
• CMSIS-NN库（ARM官方，已包含在CubeMX中）
• Python 3.x（用于权重转换脚本）

模型结构

为适配单片机资源，采用简化CNN：输入32×32×3，经2个卷积层（各含ReLU和2×2最大池化），展平后接全连接层输出10类。总参数量约12K，模型大小约48KB，推理时峰值RAM约80KB。

权重转换

使用Python脚本将训练好的Keras模型权重量化为int8并生成C头文件：

# convert_weights.py
import numpy as np
from tensorflow.keras.models import load_model

model = load_model('cifar10_cnn.h5')
# 提取权重并量化到int8
def quantize(w, scale):
    return np.clip(np.round(w / scale), -128, 127).astype(np.int8)

with open('weights.h', 'w') as f:
    f.write('#ifndef WEIGHTS_H\n#define WEIGHTS_H\n\n')
    # 卷积层1: 权重[3,3,3,16], 偏置[16]
    w1 = model.layers[0].get_weights()[0]
    b1 = model.layers[0].get_weights()[1]
    s1 = np.max(np.abs(w1)) / 127.0
    f.write(f'const int8_t conv1_w[{w1.size}] = {{' + ','.join(map(str, quantize(w1, s1).flatten())) + '};\n')
    f.write(f'const int32_t conv1_b[{b1.size}] = {{' + ','.join(map(str, b1.astype(np.int32))) + '};\n')
    f.write(f'const float conv1_scale = {s1:.8f};\n\n')
    # 卷积层2: 权重[3,3,16,32], 偏置[32]
    w2 = model.layers[2].get_weights()[0]
    b2 = model.layers[2].get_weights()[1]
    s2 = np.max(np.abs(w2)) / 127.0
    f.write(f'const int8_t conv2_w[{w2.size}] = {{' + ','.join(map(str, quantize(w2, s2).flatten())) + '};\n')
    f.write(f'const int32_t conv2_b[{b2.size}] = {{' + ','.join(map(str, b2.astype(np.int32))) + '};\n')
    f.write(f'const float conv2_scale = {s2:.8f};\n\n')
    # 全连接层: 权重[2048,10], 偏置[10]
    w3 = model.layers[4].get_weights()[0]
    b3 = model.layers[4].get_weights()[1]
    s3 = np.max(np.abs(w3)) / 127.0
    f.write(f'const int8_t fc_w[{w3.size}] = {{' + ','.join(map(str, quantize(w3, s3).flatten())) + '};\n')
    f.write(f'const int32_t fc_b[{b3.size}] = {{' + ','.join(map(str, b3.astype(np.int32))) + '};\n')
    f.write(f'const float fc_scale = {s3:.8f};\n\n')
    f.write('#endif\n')
print('weights.h generated')

单片机端实现

// main.c 核心推理代码
#include "arm_nnfunctions.h"
#include "weights.h"

// 输入图像缓冲区（32x32x3 RGB）
static int8_t img_data[32*32*3];
// 各层中间缓冲区
static int8_t conv1_out[16*16*16];
static int8_t conv2_out[8*8*32];
static int8_t fc_input[8*8*32];
static int8_t fc_out[10];

// CMSIS-NN卷积参数结构
static arm_cnn_conv_params conv1_params = {
    .padding = 1, .stride_x = 1, .stride_y = 1
};
static arm_cnn_conv_params conv2_params = {
    .padding = 1, .stride_x = 1, .stride_y = 1
};

// 量化参数（需根据实际训练调整）
static const int32_t conv1_out_shift = -8;
static const int32_t conv2_out_shift = -8;
static const int32_t fc_out_shift = -8;

void cnn_inference(int8_t *input, int8_t *output) {
    // 卷积层1 + ReLU + 最大池化
    arm_convolve_s8(input, 32*32*3, 32, 32, 3,
                    conv1_w, conv1_b, conv1_out, 16,
                    conv1_params.padding, conv1_params.stride_x,
                    conv1_params.stride_y, 0, NULL,
                    conv1_out_shift, 1, NULL, 0, NULL);
    arm_relu_q7(conv1_out, conv1_out, 16*16*16);
    arm_maxpool_s8(conv1_out, conv1_out, 16, 16, 16, 2, 2, 0, 0);

    // 卷积层2 + ReLU + 最大池化
    arm_convolve_s8(conv1_out, 16*16*16, 16, 16, 16,
                    conv2_w, conv2_b, conv2_out, 32,
                    conv2_params.padding, conv2_params.stride_x,
                    conv2_params.stride_y, 0, NULL,
                    conv2_out_shift, 1, NULL, 0, NULL);
    arm_relu_q7(conv2_out, conv2_out, 8*8*32);
    arm_maxpool_s8(conv2_out, conv2_out, 8, 8, 32, 2, 2, 0, 0);

    // 展平并全连接
    memcpy(fc_input, conv2_out, 8*8*32);
    arm_fully_connected_s8(fc_input, fc_w, fc_b, fc_out, 8*8*32, 10,
                           fc_out_shift, 1, NULL, 0, NULL);

    // 输出argmax
    int max_idx = 0;
    for (int i = 1; i < 10; i++) {
        if (fc_out[i] > fc_out[max_idx]) max_idx = i;
    }
    *output = max_idx;
}

int main(void) {
    // 系统初始化...
    // 假设img_data已填充CIFAR-10测试图像（RGB顺序，0-255映射到-128~127）
    int8_t result;
    cnn_inference(img_data, &result);
    
    // 串口输出结果
    printf("Predicted class: %d\r\n", result);
    while(1);
}

运行与验证

1. 在STM32CubeIDE中创建项目，选择STM32F407VGT6，配置USART2（115200波特率）用于输出。
2. 将生成的weights.h添加到项目Include路径。
3. 在CubeMX中启用CMSIS-NN库（Middleware → CMSIS-NN）。
4. 编译并烧录到开发板。
5. 准备一张CIFAR-10测试图像（如飞机），将其RGB数据按顺序填入img_data数组。

验证方法：串口输出"Predicted class: X"，其中X为0-9的整数。对照CIFAR-10标签（0:飞机, 1:汽车, 2:鸟, 3:猫, 4:鹿, 5:狗, 6:青蛙, 7:马, 8:船, 9:卡车），若输出与输入图像类别一致，则部署成功。可通过更换不同类别图像多次测试确认分类正确性。

常见问题

• 内存不足：若编译报RAM超限，可减小输入图像至16×16或减少卷积核数量。
• 输出恒为0：检查权重量化scale是否正确，以及输入数据是否已正确映射到int8范围。
• 推理结果错误：确认CMSIS-NN的shift参数与训练时的量化配置一致，必要时调整conv1_out_shift等参数。