在单片机上运行卷积神经网络是一种什么样的体验

这个项目是我在尝试进行表达式搜索以后开始的,目的是测试一下在MCU上跑一些小规模网络的可能性,整套开发都是基于F401CC单片机进行,最终效果可以说比较喜人,RNN和DNN的计算速度相当给力,实际的应用价值还是很高的。

概述

我在这个项目中基于DSP库的矩阵运算,手写了一套类似keras的常用层函数,包括:

函数名 说明
nn_dense 全连接层
nn_rnn 循环神经网络层,类似Keras里的SimpleRNN
nn_conv2d 2D卷积层
nn_pool2d 2D池化层
nn_flatten 展开,将深度卷积的结果展开到一维
active_softmax 多分类问题时激活用,但不等价与真正的softmax

上面的函数可以满足大多数网络的计算要求,只要单片机性能和资源足够,也可以实现一些比较有代表性的模型。

我选择的是一款STM32的单片机,具体型号为 STM32F401CC ,这个单片机拥有 64K 的RAM和 256K 的Flash,资源配置相比其他F4系列的芯片非常普通,但它依然顺利跑完了模型。

先看一下实际的运行效果,数据库来源于 MNIST-Fashion,经过Python脚本处理后,将前14张的照片转成用于OLED显示的字节数据,如果想看有关图像转换部分的实现,可以看我后续的文档介绍。

模型训练

模型的训练并不在单片机中实现,这是前提,所以并非在单片机上跑神经网络就表示既要在单片机上训练,又要在单片机上测试,这个误解不能有。我们的目的是用单片机去执行训练好的模型,而非在单片机上从头开始,试想一下,一个只有 64K RAM 的小芯片,它能做什么呢?

从这个误解中走出来,相对就方便多了,毕竟重构网络的计算比重构网络的训练要简单多了。目前大部分的深度学习网络都离不开矩阵运算,用单片机重新实现矩阵运算今儿重构网络模型就是一条非常有意义的思路。

实现方法

模型的训练在PC端完成即可,我使用了Keras来搭建模型,通过Python完成模型的训练,然后我写了一个函数用来把训练好的模型的权重信息导出。模型的权重信息本质上是一堆浮点型的数据,在Keras的模型里就是矩阵。这个矩阵会转成C语言里的数组数据,定义在 Weight.h 文件中,不过好在我用Python写了一段脚本,可以直接把模型权重转换成 weight.h 文件,省去了很多事。

实现过程

首先在PC端搭建网络模型,由于我在单片机上已经实现了一个 CNN 网络,所以我可以直接使用卷积运算,模型的结构是这样的:

模型搭建完成后即可训练,使用 Adam 作为优化器进行训练后,模型可以达到至少94%的准确率,准确率不算特别高,但对于验证单片机跑神经网络这个事件本身而言足够了。训练完成后可以导出模型的权重,通过模型结构流程图可以知道,模型的权重矩阵的shape为:

所以原则上 weight.h 需要定义 8 个数组,每一个 kernel矩阵 都需要有一个 偏置矩阵 对应,因为有四层输出,所以需要 8 个数组来定义。

好在,Python脚本帮我完成了转换。

单片机编程

STM32F4可以使用Mbed平台进行编程,我使用PlatformIO平台借助Mbed对F401编写代码。卷积神经网络需要用到卷积层、池化层、Flatten层和Softmax层,所以我提前完成了2D卷积层、2D池化层和其他需要层的定义。首先以2D卷积层为例,定义结构体如下:

// conv2d层输入结构体
typedef struct {
    uint16_t hidden_units;
    uint16_t filter_size;
    uint16_t step_size;
    uint16_t input_rows;
    uint16_t input_cols;
    uint16_t input_channel;
    conv_padding_t padding_type;
    float32_t *filter_weight;
    float32_t *filter_b;
    void (*afun)(float32_t in[], uint16_t len);
} conv2d_input_t;

2D池化层定义如下:

// pool2d层输入结构体,MaxPooling2D AvgPooling2D 可用
typedef struct {
    uint16_t input_rows;
    uint16_t input_cols;
    uint16_t input_channels;
    uint8_t pooling_size;
    uint8_t pooling_step;
    float32_t (*poolng_type_fun)(float32_t in[], uint16_t len);
} pool2d_input_t;

FLatten层的定义如下:

// flatten展开结构体
typedef struct {
    uint16_t input_rows;
    uint16_t input_cols;
    uint16_t input_channels;
} flatten_input_t;

用到的函数原型定义:

// 池化辅助函数
float32_t nn_pooling_max(float32_t data[], uint16_t len);
float32_t nn_pooling_min(float32_t data[], uint16_t len);
float32_t nn_pooling_avg(float32_t data[], uint16_t len);
// 激活函数
void active_relu(float32_t in[], uint16_t len);
void active_none(float32_t in[], uint16_t len);
uint16_t active_softmax(float32_t res[], uint16_t len);
// 网络结构
void nn_dense(dense_input_t inp_struct, float32_t input[], float32_t output[]);
void nn_conv2d(conv2d_input_t inp_struct, float32_t input[], float32_t output[]);
void nn_pool2d(pool2d_input_t inp_struct, float32_t input[], float32_t output[]);
void nn_flatten(flatten_input_t inp_struct, float32_t input[], float32_t output[]);

需要说明一下的是,这里的softmax并不是真正意义上的softmax函数,因真正的softmax函数需要根据输出记性概率转化,这个过程计算起来非常复杂,单片机的能力有限,所以这里的softmax其实是取最后一个输出层的最大值索引,因为最大值往往通过softmax计算之后也对应了最大的概率,所以最终结果是一致。

最后贴上实现整个模型的C代码:

// CNN模型的实现
uint8_t run_model(float32_t data[])
{
    // 卷积结构体
    conv2d_input_t conv2d_1_inp_data;
    conv2d_1_inp_data.afun = &active_relu;
    conv2d_1_inp_data.filter_size = CONV_FILTER_3X3;
    conv2d_1_inp_data.filter_weight = conv2d_1_w;
    conv2d_1_inp_data.filter_b = conv2d_1_b;
    conv2d_1_inp_data.step_size = CONV_FILTER_STEP_1;
    conv2d_1_inp_data.input_rows = C_INPUT_WIDTH;
    conv2d_1_inp_data.input_cols = C_INPUT_HEIGHT;
    conv2d_1_inp_data.input_channel = 1;
    conv2d_1_inp_data.padding_type = CONV_PADDING_NONE;
    conv2d_1_inp_data.hidden_units = C_CONV2D_1_UNITS;
    // 输入值
    float32_t *input = data;
    // 输出值
    uint8_t padding = conv2d_1_inp_data.padding_type == CONV_PADDING_NONE ? 2 : 0;
    uint8_t out_1_width = conv2d_1_inp_data.input_rows - padding;
    uint8_t out_1_height = conv2d_1_inp_data.input_cols - padding;
    uint16_t array_size = (uint16_t)(out_1_width * out_1_height / conv2d_1_inp_data.step_size);
    float32_t *output_1 = (float32_t *)malloc(sizeof(float32_t) * array_size * C_CONV2D_1_UNITS);
    // 运行conv2d层
    nn_conv2d(conv2d_1_inp_data, input, output_1);

    // 第二层卷积
    conv2d_input_t conv2d_2_inp_data;
    conv2d_2_inp_data.afun = &active_relu;
    conv2d_2_inp_data.filter_size = CONV_FILTER_3X3;
    conv2d_2_inp_data.filter_weight = conv2d_2_w;
    conv2d_2_inp_data.filter_b = conv2d_2_b;
    conv2d_2_inp_data.step_size = CONV_FILTER_STEP_1;
    conv2d_2_inp_data.input_rows = out_1_width;
    conv2d_2_inp_data.input_cols = out_1_height;
    conv2d_2_inp_data.input_channel = C_CONV2D_1_UNITS;
    conv2d_2_inp_data.padding_type = CONV_PADDING_NONE;
    conv2d_2_inp_data.hidden_units = C_CONV2D_2_UNITS;
    // 输出值
    padding = conv2d_2_inp_data.padding_type == CONV_PADDING_NONE ? 2 : 0;
    uint8_t out_2_width = conv2d_2_inp_data.input_rows - padding;
    uint8_t out_2_height = conv2d_2_inp_data.input_cols - padding;
    array_size = (uint16_t)(out_2_width * out_2_height / conv2d_2_inp_data.step_size);
    float32_t *output_2 = (float32_t *)malloc(sizeof(float32_t) * array_size * C_CONV2D_2_UNITS);
    // 运行conv2d层
    nn_conv2d(conv2d_2_inp_data, output_1, output_2);
    free(output_1);

    // 第一层池化层
    pool2d_input_t pool_1_inp_data;
    pool_1_inp_data.input_channels = C_CONV2D_2_UNITS;
    pool_1_inp_data.input_rows = out_2_width;
    pool_1_inp_data.input_cols = out_2_height;
    pool_1_inp_data.pooling_size = POOLING_SIZE_2X2;
    pool_1_inp_data.pooling_step = POOLING_STEP_2;
    pool_1_inp_data.poolng_type_fun = &nn_pooling_max;
    // 输出值
    uint16_t poolout_1_width = out_2_width / POOLING_SIZE_2X2;
    uint16_t poolout_1_height = out_2_height / POOLING_SIZE_2X2;
    array_size = (uint16_t)(poolout_1_width * poolout_1_height);
    float32_t *pool_output_1 = (float32_t *)malloc(sizeof(float32_t) * array_size * C_CONV2D_2_UNITS);
    // 计算池化
    nn_pool2d(pool_1_inp_data, output_2, pool_output_1);
    free(output_2);

    // 第三层卷积
    conv2d_input_t conv2d_3_inp_data;
    conv2d_3_inp_data.afun = &active_relu;
    conv2d_3_inp_data.filter_size = CONV_FILTER_3X3;
    conv2d_3_inp_data.filter_weight = conv2d_3_w;
    conv2d_3_inp_data.filter_b = conv2d_3_b;
    conv2d_3_inp_data.step_size = CONV_FILTER_STEP_1;
    conv2d_3_inp_data.input_rows = poolout_1_width;
    conv2d_3_inp_data.input_cols = poolout_1_height;
    conv2d_3_inp_data.input_channel = C_CONV2D_2_UNITS;
    conv2d_3_inp_data.padding_type = CONV_PADDING_NONE;
    conv2d_3_inp_data.hidden_units = C_CONV2D_3_UNITS;
    // 输出值
    padding = conv2d_3_inp_data.padding_type == CONV_PADDING_NONE ? 2 : 0;
    uint8_t out_3_width = conv2d_3_inp_data.input_rows - padding;
    uint8_t out_3_height = conv2d_3_inp_data.input_cols - padding;
    array_size = (uint16_t)(out_3_width * out_3_height / conv2d_3_inp_data.step_size);
    float32_t *output_3 = (float32_t *)malloc(sizeof(float32_t) * array_size * C_CONV2D_3_UNITS);
    // 运行conv2d层
    nn_conv2d(conv2d_3_inp_data, pool_output_1, output_3);
    free(pool_output_1);

    // 第二池化层
    pool2d_input_t pool_2_inp_data;
    pool_2_inp_data.input_channels = C_CONV2D_3_UNITS;
    pool_2_inp_data.input_rows = out_3_width;
    pool_2_inp_data.input_cols = out_3_height;
    pool_2_inp_data.pooling_size = POOLING_SIZE_2X2;
    pool_2_inp_data.pooling_step = POOLING_STEP_2;
    pool_2_inp_data.poolng_type_fun = &nn_pooling_max;
    // 输出值
    uint16_t poolout_2_width = out_3_width / POOLING_SIZE_2X2;
    uint16_t poolout_2_height = out_3_height / POOLING_SIZE_2X2;
    array_size = (uint16_t)(poolout_2_width * poolout_2_height);
    float32_t *pool_output_2 = (float32_t *)malloc(sizeof(float32_t) * array_size * C_CONV2D_3_UNITS);
    // 计算池化
    nn_pool2d(pool_2_inp_data, output_3, pool_output_2);
    free(output_3);

    // Flatten层
    flatten_input_t flatten_inp_data;
    flatten_inp_data.input_rows = poolout_2_width;
    flatten_inp_data.input_cols = poolout_2_height;
    flatten_inp_data.input_channels = C_CONV2D_3_UNITS;
    // 输出值
    uint16_t outsize = flatten_inp_data.input_rows * flatten_inp_data.input_cols * flatten_inp_data.input_channels;
    float32_t *flatten_output = (float32_t *)malloc(sizeof(float32_t) * outsize);
    // 展开
    nn_flatten(flatten_inp_data, pool_output_2, flatten_output);
    free(pool_output_2);

    // 第一全连接层
    arm_matrix_instance_f32 mat_dense1_w = {C_FALTTEN_UNITS, C_DENSE_1_UNITS, dense_1_w};
    arm_matrix_instance_f32 mat_dense1_b = {1, C_DENSE_1_UNITS, dense_1_b};
    // 输出
    float32_t *dense1_out = (float32_t *)malloc(sizeof(float32_t) * C_DENSE_1_UNITS);
    dense_input_t dense_1_inp_data;
    dense_1_inp_data.afun = &active_none;
    dense_1_inp_data.input_units = C_FALTTEN_UNITS;
    dense_1_inp_data.dense_units = C_DENSE_1_UNITS;
    dense_1_inp_data.mat_weight = &mat_dense1_w;
    dense_1_inp_data.mat_b = &mat_dense1_b;
    // 计算全连接层
    nn_dense(dense_1_inp_data, flatten_output, dense1_out);
    free(flatten_output);

    // 最后一层,softmax激活
    uint8_t res_index = active_softmax(dense1_out, C_DENSE_1_UNITS);
    free(dense1_out);

    // 返回分类索引
    return res_index;
}

之前使用循环神经网络跑MNIST数据集时可以达到每张图6ms的速度,可以说非常快了,在使用这个卷积神经网络进行识别分类时,可以达到10ms,也就是10帧的速度,速度也还可以的,毕竟,MNIST的输入图只有 28×28 大小。

总结

总体来说,单片机的运行资源比较有限,想实现高难度大规模的模型是不太可能的,但是另一方面,运行一些优化好且体积较小的模型也是可行的,这里主要依赖于单片机的资源条件。总体来说,单片机的运行频率决定了模型能跑多快,单片机的片内资源大小决定了模型能跑多大。有关这部分的总结,我写在了下一篇博文中,对我在开发过程中遇到的问题,已经开发经验做了整理,或许会有一些参考的价值。

类似文章

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注