在STM32上实现嵌入式人工智能应用

news2024/11/25 6:51:39

引言

随着微控制器的计算能力不断增强，人工智能（AI）开始在嵌入式系统中扮演越来越重要的角色。STM32微控制器由于其高性能和低功耗的特性，非常适合部署轻量级AI模型。

本文将探讨如何在STM32平台上实现深度学习应用，特别是利用STM32Cube.AI工具链将训练好的神经网络模型部署到STM32设备上。

环境准备

硬件选择：STM32F746G Discovery kit，具备足够的计算资源和内存支持复杂模型。
软件需求：
- STM32CubeMX：用于配置微控制器的外设和中间件。
- STM32CubeIDE：集成开发环境，支持代码编辑、编译、调试。
- STM32Cube.AI：将神经网络模型转换为优化的代码，以在STM32设备上运行。
- TensorFlow Lite：用于训练神经网络模型。

AI模型的训练和转换

训练一个简单的神经网络

使用TensorFlow Lite训练一个用于分类任务的简单神经网络模型。此示例中，我们将训练一个模型来识别手写数字（基于MNIST数据集）。

代码示例：训练模型

import tensorflow as tf

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=10)
model.evaluate(x_test, y_test)

转换模型为STM32Cube.AI格式

使用STM32Cube.AI将训练好的TensorFlow Lite模型转换为C代码，以便在STM32设备上运行。

操作步骤：

使用STM32CubeMX创建一个新项目，并配置所需的外设。
在STM32Cube.AI中导入TensorFlow模型。
生成优化后的代码，并集成到STM32项目中。

在STM32上部署和执行AI模型

初始化AI库和模型

在STM32CubeIDE中设置项目，初始化AI库，并加载模型。

代码示例：初始化AI模型

#include "app_x-cube-ai.h"
#include "ai_datatypes_defines.h"

extern AI_ALIGNED(4) ai_u8 activations[AI_NETWORK_DATA_ACTIVATIONS_SIZE];

ai_handle network = AI_HANDLE_NULL;

void aiInit(void) {
    ai_error err;

    const ai_network_params ai_params = {
        AI_NETWORK_DATA_WEIGHTS(ai_network_data_weights_get()),
        AI_NETWORK_DATA_ACTIVATIONS(activations)
    };

    err = ai_network_create(&network, AI_NETWORK_DATA_CONFIG);
    if (err.type != AI_ERROR_NONE) {
        // Handle error
    }
    if (!ai_network_init(network, &ai_params)) {
        // Handle error
    }
}

执行推理

在STM32上执行AI模型的推理，处理输入数据并获取预测结果。

代码示例：执行推理

void aiRun(float* input, float* output) {
    ai_i32 nbatch;
    ai_buffer ai_input[1];
    ai_buffer ai_output[1];

    // Configure input buffer
    ai_input[0] = ai_network_inputs_get(network, NULL);
    ai_input[0].data = AI_HANDLE_PTR(input);

    // Configure output buffer
    ai_output[0] = ainetwork_outputs_get(network, NULL);
    ai_output[0].data = AI_HANDLE_PTR(output);
    // Perform the inference
    nbatch = ai_network_run(network, ai_input, ai_output);
    if (nbatch != 1) {
    // Handle error
}

应用场景