极客秀
搜索

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

模型文件和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:SYSTEMCurrentControlSetControlFileSystem" `-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函数,他有两个参数,一个是输入数据数组,一个是输出数据数组,由于我训练模型的时候将图片大小调整为4040,所以是一个4040的输入,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);        }      }    }  }

我们将屏幕范围内400400的区域划分成4040个像素块,记录下我们的手指划过轨迹存入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(); // 清除所有点亮的方块  }

接下来当我们按下按键的时候,将In_data数组丢入神经网络训练,输出的结果在out_data中,我们从out_data中挑选出置信度最高的值输出就完成了我们的手写数字识别的功能了。

** 总结 **
总的来说使用下来几乎没有什么卡顿,之前一直是低于CubeAI的强大了,结合那位工程师的指教,在视觉方面即便自己的代码能力欠缺也可以从开源模型库中挑选优化好的适合部署以及满足要求的模型直接利用CubeAI来使用,使用起来也非常方便。
自己也非常感慨,因为去年去研究的时候压根没研究明白,一点头绪也没有,现在好歹能和去年的自己构成了闭环,也是勉励自己一直学习下去。

1.转载请保留原文链接谢谢!
2.本站所有资源文章出自互联网收集整理,本站不参与制作,如果侵犯了您的合法权益,请联系本站我们会及时删除。
3.本站发布资源来源于互联网,可能存在水印或者引流等信息,请用户擦亮眼睛自行鉴别,做一个有主见和判断力的用户。
4.本站资源仅供研究、学习交流之用,若使用商业用途,请购买正版授权,否则产生的一切后果将由下载用户自行承担。
5.联系方式(#替换成@):pm#vimge.com

  相关内容