如何在STM32中利用CubeAI部署自己的神经网络模型（以手写数字识别为例）

redparrot2008

已于 2025-02-24 15:00:26 修改

阅读量2.1k

点赞数 12

分类专栏：神经网络文章标签： stm32 神经网络嵌入式硬件

于 2025-01-08 17:25:08 首次发布

原文链接：https://mp.weixin.qq.com/s/Nzyf5vrMKZtCJVORw_8dBg

版权

神经网络专栏收录该内容

3 篇文章

订阅专栏

CubeMx 和 STM32Cube.AI

模型文件和STM32的工程有需要可以私信：手写识别例程 ，开发板用的是正点原子的F407探索者开发板。

STM32Cube.AI可用来在任意STM32微控制器上,对采用最流行AI框架训练的神经网络模型进行优化和部署。它是一款可以集成在CubeMX（图形初始化工具）中可以在初始化时部署AI模型的插件。

我在去年第一次了解到CubeAI的时候就对其产生了浓厚的兴趣，思考着在STM32中如何快速部署神经网络模型，但是由于自己不是这方面专业的，对神经网络训练等一窍不通，因此一直没有时间去实验。

后来ST推出了NanoEdge来帮我们这些门外汉快速的部署神经网络模型，但是它的局限性比较大，不太适合对图像等进行训练，因此例如手写识别，图像识别等还是需要依靠CubeAI插件将一个训练好的神经网络模型嵌入到单片机中。在九月份有幸和ST相关方面的技术人员聊了这部分内容关于NanoEdge的局限性，他建议我可以从开源模型库中选择合适的模型来导入而不一定要自己训练和调节参数，这确实是一个方法，不过得益于AI技术的发展，我们对于自己不擅长的领域可以得到很大的帮助，话不多说，我们来介绍一下从模型训练到部署的过程。

模型训练

def create_mnist_model():    print("TensorFlow版本:", tf.__version__)
    # 加载MNIST数据集    (train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
    # 数据预处理：将图像调整为40x40    def resize_images(images):        resized = []        for img in images:            img = tf.image.resize(img[..., tf.newaxis], [40, 40]).numpy()            resized.append(img)        return np.array(resized)
    train_images = resize_images(train_images)    test_images = resize_images(test_images)
    train_images = train_images.astype('float32') / 255    test_images = test_images.astype('float32') / 255
    # 创建模型    model = models.Sequential([        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(40, 40, 1)),        layers.MaxPooling2D((2, 2)),        layers.Conv2D(64, (3, 3), activation='relu'),        layers.MaxPooling2D((2, 2)),        layers.Conv2D(64, (3, 3), activation='relu'),        layers.Flatten(),        layers.Dense(64, activation='relu'),        layers.Dense(10, activation='softmax')    ])
    model.compile(optimizer='adam',                  loss='sparse_categorical_crossentropy',                  metrics=['accuracy'])
    # 训练模型    print("开始训练模型...")    history = model.fit(train_images, train_labels, epochs=5, batch_size=32,              validation_split=0.2, verbose=1)
    # 评估模型    test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)    print(f"测试准确率: {test_acc:.4f}")
    return model, history, test_acc

这是一段Python代码，我主要采用了Tensorflow来训练模型，数据集来自Tensorflow中的keras模块中加载MNIST数据集，只有将其轻量化处理，保存为.tflite文件，注意不是.tf格式而是轻量化后的.tflite

因为CubeAI好像只支持这三种，模型量化前后的大小会差距很大，我的F407只有1MB的flash，量化后的模型大小是600KB，单片机完全存放的下。

Keil以及CubeMX

之后打开CubeMX创建一个单片机的工程，安装CubeAI插件并激活。

在CubeAI配置界面导入我们的模型。

之后验证我们模型是否可以被部署。

这里出错可能是文件路径太长了，可以用管理员模式打开PowerShell用如下代码开启长路径模式，然后重启电脑就可以解决了。

New-ItemProperty -Path "HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem" `-Name "LongPathsEnabled" -Value 1 -PropertyType DWORD -Force

分析完成后就可以生成我们的工程了，当然前面的外设初始化也要正常的完成。

可以看到工程中出现了神经网络的相关文件还有模型文件。

接下来我们来介绍如何使用。

static ai_handle network = AI_HANDLE_NULL;

/* Global c-array to handle the activations buffer *//*ڴȫc*/
AI_ALIGNED(32)
static ai_u8 activations[AI_NETWORK_DATA_ACTIVATIONS_SIZE];

/* Array to store the data of the input tensor *//*ڴ洢?*/
AI_ALIGNED(32)
static ai_float in_data[AI_NETWORK_IN_1_SIZE]; // 这里假设 AI_NETWORK_IN_1_SIZE 为 1600 (40x40)
/* or static ai_u8 in_data[AI_NETWORK_IN_1_SIZE_BYTES]; *//*ڴ洢ݵc*/

/* c-array to store the data of the output tensor */
AI_ALIGNED(32)
static ai_float out_data[AI_NETWORK_OUT_1_SIZE];
/* static ai_u8 out_data[AI_NETWORK_OUT_1_SIZE_BYTES]; */

/* Array of pointer to manage the model's input/output tensors */
/*ڹģ/ָ*/
static ai_buffer *ai_input;
static ai_buffer *ai_output;

/* 
 * Bootstrap
 */
int aiInit(void) {
  ai_error err;
  
  /* Create and initialize the c-model *//*ʼc-model */
  const ai_handle acts[] = { activations };
  err = ai_network_create_and_init(&network, acts, NULL);
  if (err.type != AI_ERROR_NONE) {  };

  /* Reteive pointers to the model's input/output tensors *//*ȡָģ/ָ*/
  ai_input = ai_network_inputs_get(network, NULL);
  ai_output = ai_network_outputs_get(network, NULL);

  return 0;
}

/* 
 * Run inference
 */
int aiRun(const void *in_data, void *out_data) {
  ai_i32 n_batch;
  ai_error err;
  
  /* 1 - Update IO handlers with the data payload *//* 1 -ЧغɸIO*/
  ai_input[0].data = AI_HANDLE_PTR(in_data);
  ai_output[0].data = AI_HANDLE_PTR(out_data);

  /* 2 - Perform the inference *//* 2 -ִ*/
  n_batch = ai_network_run(network, &ai_input[0], &ai_output[0]);
  if (n_batch != 1) {
    err = ai_network_get_error(network);
    HAL_GPIO_TogglePin(GPIOF,GPIO_PIN_10);
  };
  
  return 0;
}

这里我们引用一下官网的参考代码，主要有两个部分，一个是aiInit用于输出化这个可以直接调用，其次就是aiRun函数，他有两个参数，一个是输入数据数组，一个是输出数据数组，由于我训练模型的时候将图片大小调整为40*40，所以是一个40*40的输入，10输出（数字0~9）的置信。

  while(1)  {    tp_dev.scan(1);    if(tp_dev.sta == 0xC1)    {      uint16_t x = tp_dev.x[0];      uint16_t y = tp_dev.y[0];            // 计算触摸点所在的网格位置      int gridX = x / 10;      int gridY = y / 10;            // 标记该网格为已触摸，并点亮对应的10x10方块      if(gridX < 40 && gridY < 40) {        int index = gridY * 40 + gridX;        if(in_data[index] == 0) {          in_data[index] = 1.0f;          LCD_Fill(gridX*10, gridY*10, (gridX+1)*10-1, (gridY+1)*10-1, RED);        }      }    }  }

我们将屏幕范围内400*400的区域划分成40*40个像素块，记录下我们的手指划过轨迹存入in_data数组中。

KeyCallBack(&key,GPIOE,GPIO_PIN_2);if(key.HasPressNum){  HAL_GPIO_TogglePin(GPIOF,GPIO_PIN_9);  key.HasPressNum = 0;  aiRun(in_data,out_data);  // 找出 out_data 中的最大值及其索引  int max_index = 0;  float max_value = out_data[0];  for (int i = 1; i < AI_NETWORK_OUT_1_SIZE; i++) {      if (out_data[i] > max_value) {          max_value = out_data[i];          max_index = i;        }   }          // 在这里可以使用 max_index 和 max_value，将结果显示在LCD上或通过其他方式输出          char result[20];  sprintf(result, "The most like: %d (%.2f)", max_index, max_value);  LCD_ShowString(10, 410, 400, 32, 32, (uint8_t*)result);  clearAllTouchedBlocks(); // 清除所有点亮的方块  }