60行代码加速20倍 NEON实现深度学习OD任务后处理绘框

前言由于此前OD任务发现在检测后处理部分使用OpenCV较为占用资源且耗时，本文遂尝试使用NEON做后处理绘框，以达到加速并降低CPU资源消耗的目的。
Pytorch训练营，花两个星期彻底掌握代码实现
CV各大方向专栏与各个部署框架最全教程整理
CV全栈指导班、基础入门班、论文指导班全面上线!!

一、实现思路

假设对一张Mat图像进行操作（其实也不仅仅是Mat对象，理论上只要知道图像通道的首指针即可），在ARM端使用NEON instrinc指令集里实现一个后处理绘框的功能，可以简单罗列成以下几步：1. 定义参数： 首先确定图像的宽度和高度，图像的首地址指针，以及边界（边框）的厚度。2. 向量寄存器加载： 使用NEON的加载指令从内存中加载像素数据到向量寄存器中。3. 处理上下边框：

对于顶部边界，遍历整个第一行的像素，并使用NEON的存储指令将特定颜色值写回到这些位置（比如想绘制的是绿框，那么需要将B通道的绘框元素数据更改为0，G通道为255，R通道为0）。
同样地，对于底部边界，遍历最后一行的像素并执行相同的操作。

4.处理左右边框：这个稍微复杂一些，因为需要处理每一行的开始和结束位置。一种方法是使用循环，每次处理一行，然后更新寄存器中的值以反映特定颜色。我们可以使用NEON的广播指令来创建一个包含特定颜色所有分量的向量，然后使用存储指令将其写入到图像的左侧和右侧边界。5.边框优化：由于很多检测框的宽度很难保证一定是SIMD WIDTH的倍数，这就造成了在绘图时一些不必要的麻烦，举个例子，假设检测框的width是97，SIMD WIDTH的长度是16（一次性处理16个元素的向量寄存器），那么97/16=6······1，刚好多出了1个pixel，此时需要某些处理措施规避这种情况。

二、实现过程

2.1 定义参数

首先确定图像的宽度和高度，本次测试所获得的检测框均由这篇博文中的end2end模型中获得【1】，也就是在绘框前，我们会得到一个vector数组，均为通过nms获得的检测框，这个数组数据排列格式如下：一个box对应四个元素，其实box是按照obj的score排列，但为了方便讲解，我们假设他是按从左到右顺序排列，由于测试的图片均为COCO2017 Val中的数据，图片尺寸中值远大于320，为了美观，此篇博文默认绘框边界（边框）的厚度为2，也就是占满2个pixel。函数定义如下：

void neon_rectangle_blod(uint8_t *img, uint16_t img_step, uint16_t x, uint16_t y, uint16_t w, uint16_t h, uint8_t blue, uint8_t green, uint8_t red)

函数形参解释：为图像首指针，指图像的width，指检测框左上角，指检测框的宽高，指三色通道需要填充的数值。

2.2 向量寄存器加载

这一步需要将图像BGR通道元素加载到寄存器，由于图像一般为uint8格式，这里可以使用最大的寄存器，把位宽拉满，也就是一次性操作16个元素，调用NEON instrinc中的vld3q_u8加载图像BGR数据到uint8x16x3_t寄存器中，再将单个通道的数据分发到到单个uint8x16_t寄存器中，伪代码如下:

// 假设img是指向图像BGR数据的指针
uint8x16x3_t bgr_data = vld3q_u8((uint8_t *) img);

// 分别将BGR通道的数据分发到单独的uint8x16_t寄存器
uint8x16_t reg_b = bgr_data.val[0]; // 蓝色通道
uint8x16_t reg_g = bgr_data.val[1]; // 绿色通道
uint8x16_t reg_r = bgr_data.val[2]; // 红色通道

// 后续对每个通道进行单独的操作
......

2.3 处理上下边框

我们需要定位到上下边框的起始位置，获取起始位置的地址，再将地址往后以16个pixel为一个SIMD_WIDTH塞入寄存器，将寄存器中的B,G,R通道进行向量赋值，表示一次性处理16个数据流位宽，代码如下:

// 绘制矩形的上下边界
    for (uint16_t i = 0; i 16)
    {
        // 计算当前行的起始地址
        uint8_t *top_row1 = img + (y * img_step + x + i) * 3;
        uint8_t *bottom_row1 = img + ((y + h) * img_step + x + i) * 3;

        // 使用NEON指令集并行加载和存储颜色
        uint8x16x3_t pixels_top1 = vld3q_u8(top_row1);
        uint8x16x3_t pixels_bottom1 = vld3q_u8(bottom_row1);

        // 绘制顶部和底部线条
        pixels_top1.val[0] = neon_color_b; // 蓝色通道
        pixels_top1.val[1] = neon_color_g; // 绿色通道
        pixels_top1.val[2] = neon_color_r; // 红色通道


    


        pixels_bottom1.val[0] = neon_color_b;
        pixels_bottom1.val[1] = neon_color_g;
        pixels_bottom1.val[2] = neon_color_r;

        vst3q_u8(top_row1, pixels_top1);
        vst3q_u8(bottom_row1, pixels_bottom1);

        // 计算当前行的起始地址
        uint8_t *top_row2 = img + ((y + 1) * img_step + x + i) * 3;
        uint8_t *bottom_row2 = img + ((y + h - 1) * img_step + x + i) * 3;

        // 使用NEON指令集并行加载和存储颜色
        uint8x16x3_t pixels_top2 = vld3q_u8(top_row2);
        uint8x16x3_t pixels_bottom2 = vld3q_u8(bottom_row2);

        // 绘制顶部和底部线条
        pixels_top2.val[0] = neon_color_b; // 蓝色通道
        pixels_top2.val[1] = neon_color_g; // 绿色通道
        pixels_top2.val[2] = neon_color_r; // 红色通道

        pixels_bottom2.val[0] = neon_color_b;
        pixels_bottom2.val[1] = neon_color_g;
        pixels_bottom2.val[2] = neon_color_r;

        vst3q_u8(top_row2, pixels_top2);
        vst3q_u8(bottom_row2, pixels_bottom2);
    }

2.4 处理左右边框

这里就有点难受了，因为是ARM架构通用的汇编，不像一些厂家有专门处理竖直方向的寄存器或者额外的硬件加速模块，所以这一步只能老老实实一个pixel一个pixel的去涂，因此和OpenCV的处理方式没有太大差异，代码如下:

// 绘制矩形的左右边界
    for (uint16_t j = 0; j     {
        // 计算当前列的起始地址
        uint8_t *left_col1 = img + ((y + j) * img_step + x) * 3;
        uint8_t *right_col1 = img + ((y + j) * img_step + (x + w)) * 3;
        // 设置左边和右边列的颜色
        left_col1[0] = right_col1[0] = blue;
        left_col1[1] = right_col1[1] = green;
        left_col1[2] = right_col1[2] = red;

        // 计算当前列的起始地址
        uint8_t *left_col2 = img + ((y + j) * img_step + x + 1) * 3;
        uint8_t *right_col2 = img + ((y + j) * img_step + (x + w) - 1) * 3;
        // 设置左边和右边列的颜色
        left_col2[0] = right_col2[0] = blue;
        left_col2[1] = right_col2[1] = green;
        left_col2[2] = right_col2[2] = red;
    }

2.5 优化边框

这里提供一种思路，既然没办法确保检测框的宽度刚好是SIMD_WIDTH的倍数，那我们就将宽度扩充或者减小到SIMD_WIDTH的倍数，但为了美观处理，不管是扩充还是减小宽度，我们都离不开一个操作，那就是中心对齐，以扩宽为例，如下图所示：那么，就有很好的方式去应对这种情况，我们假设检测框的width对SIMD_WIDTH进行mod操作，如果余数小于，对检测框width进行缩小操作，反之，则进行扩充操作，代码如下：

void check_point(int *x1, int *x2, 


    
int nstride)
{
    int mod, w, xc, nw;
    w = *x2 - *x1;
    xc = *x1 + (int)(w / 2);
    mod = w % nstride;
    if (mod > (nstride / 2))
    {
        *x1 = xc - (int)((w + nstride - mod) / 2);
        *x2 = xc + (int)((w + nstride - mod) / 2);
    }
    else
    {
        nw = w - mod;
        *x1 = xc - int(nw / 2);
        *x2 = xc + int(nw / 2);
    }
}

三、测试结果

测试机器为4+32内存的树莓派4B，共带有4颗A72核，我们分别使用NEON和OpenCV作为【1】中end2end模型出框后的后处理绘框函数，测试数据为COCO2017 Val数据集，将两个程序用taskset -c先绑定在编号为0的核上，得出两者在处理5000张图的处理速度差异，如下所示:其中，cost time为推理完5000张图的所有耗时，单位为ms，average cost time为处理单张图片的耗时，单位为us，我们可以看到，在单个A72上，NEON实现的绘框函数要比OpenCV快了20倍左右。此外，OpenCV的强大源于多核并行，为了能更加客观且全面的测试出两者的性能差异，我们在OpenCV版本的基础上，不断增加核进行测试，得出以下测试图例：图中P/ms表示1ms能处理多少图，越高表示每毫秒处理图越多，单图绘框速度越快，从图可以看出，单核运行的NEON绘框的速度依旧稳稳碾压多核并行的OpenCV。OpenCV绘框效果如下:NEON汇编绘框效果如下:

四、完整代码

void check_point(int *x1, int *x2, int nstride)
{
    int mod, w, xc, nw;
    w = *x2 - *x1;
    xc = *x1 + (int)(w / 2);
    mod = w % nstride;
    if (mod > (nstride / 2))
    {
        *x1 = xc - (int)((w + nstride - mod) / 2);
        *x2 = xc + (int)((w + nstride - mod) / 2);
    }
    else
    {
        nw = w - mod;
        *x1 = xc - int(nw / 2);
        *x2 = xc + int(nw / 2);
    }
}

void neon_rectangle_blod(uint8_t *img, uint16_t img_step, uint16_t x, uint16_t y, uint16_t w, uint16_t h, uint8_t blue, uint8_t green, uint8_t red)
{



    
    // 创建一个全1的8位向量，用于绘制矩形的颜色
    uint8x16_t neon_color_b = vdupq_n_u8(blue);
    uint8x16_t neon_color_g = vdupq_n_u8(green);
    uint8x16_t neon_color_r = vdupq_n_u8(red);

    // 绘制矩形的上下边界
    for (uint16_t i = 0; i 16)
    {
        // 计算当前行的起始地址
        uint8_t *top_row1 = img + (y * img_step + x + i) * 3;
        uint8_t *bottom_row1 = img + ((y + h) * img_step + x + i) * 3;

        // 使用NEON指令集并行加载和存储颜色
        uint8x16x3_t pixels_top1 = vld3q_u8(top_row1);
        uint8x16x3_t pixels_bottom1 = vld3q_u8(bottom_row1);

        // 绘制顶部和底部线条
        pixels_top1.val[0] = neon_color_b; // 蓝色通道
        pixels_top1.val[1] = neon_color_g; // 绿色通道
        pixels_top1.val[2] = neon_color_r; // 红色通道

        pixels_bottom1.val[0] = neon_color_b;
        pixels_bottom1.val[1] = neon_color_g;
        pixels_bottom1.val[2] = neon_color_r;

        vst3q_u8(top_row1, pixels_top1);
        vst3q_u8(bottom_row1, pixels_bottom1);

        // 计算当前行的起始地址
        uint8_t *top_row2 = img + ((y + 1) * img_step + x + i) * 3;
        uint8_t *bottom_row2 = img + ((y + h - 1) * img_step + x + i) * 3;

        // 使用NEON指令集并行加载和存储颜色
        uint8x16x3_t pixels_top2 = vld3q_u8(top_row2);
        uint8x16x3_t pixels_bottom2 = vld3q_u8(bottom_row2);

        // 绘制顶部和底部线条
        pixels_top2.val[0] = neon_color_b; // 蓝色通道
        pixels_top2.val[1] = neon_color_g; // 绿色通道
        pixels_top2.val[2] = neon_color_r; // 红色通道

        pixels_bottom2.val[0] = neon_color_b;
        pixels_bottom2.val[1] = neon_color_g;
        pixels_bottom2.val[2] = neon_color_r;

        vst3q_u8(top_row2, pixels_top2);
        vst3q_u8(bottom_row2, pixels_bottom2);
    }

    // 绘制矩形的左右边界
    for (uint16_t j = 0; j     {
        // 计算当前列的起始地址
        uint8_t *left_col1 = img + ((y + j) * img_step + x) * 3;
        uint8_t *right_col1 = img + ((y + j) * img_step + (x + w)) * 3;
        // 设置左边和右边列的颜色
        left_col1[0] = right_col1[0] = blue;
        left_col1[1] = right_col1[1] = green;
        left_col1[2] = right_col1[2] = red;

        // 计算当前列的起始地址
        uint8_t *left_col2 = img + ((y + j) * img_step + x + 1) * 3;
        uint8_t *right_col2 = img + ((y + j) * img_step + (x + w) - 1) * 3;
        // 设置左边和右边列的颜色
        left_col2[0] = right_col2[0] = blue;
        left_col2[1] = right_col2[1


    
] = green;
        left_col2[2] = right_col2[2] = red;
    }
}

五、总结

本篇博文主要讲述后处理绘框的汇编实现方式，在树莓派上的单核以及多核A72上都实现了加速，但时间关系未于其他开发板做比较，从去年开始，似乎4大+4小变成了业界主流，既4颗A76+4颗A57或者4颗A76+4颗A53，ARM端CPU算力要远远强过四颗A72，至于这种汇编实现方式，在这些开发板上能加速多少，确实不好说，有兴趣的朋友可以用这几十行代码去测试下~

六、参考

[1] https://zhuanlan.zhihu.com/p/672633849 [2] https://zhuanlan.zhihu.com/p/698551682 [3] https://developer.arm.com/documentation/

若觉得还不错的话，请点个 “赞” 或 “在看” 吧

论文指导班

论文指导班面向那些没有导师指导、需要升学申博的朋友，指导学员从零开始调研相关方向研究、尝试idea、做实验、写论文，指导老师会提供一些idea、代码实现部分的指导、论文写作指导和修改，但整体仍然是由学员自主完成。需要说明的是，论文指导班并非帮你写论文，或者直接给一篇论文让你挂名，我们不会做任何灰色产业，因此，想直接买论文或挂名的朋友请勿联系。

指导老师：

海外QS Top-60某高校人工智能科学博士在读, 师从IEEE Fellow，曾在多家AI企业担任研究实习生和全职算法研究员，具备极强的学术届和工业界综合背景。研究领域主要包括通用计算机视觉模型的高效设计，训练，部署压缩以及在目标检测，语义分割等下游任务应用，具体包括模型压缩 (知识蒸馏，模型搜索量化剪枝), 通用视觉模型与应用(VIT, 目标检测，语义分割), AI基础理论(AutoML, 数据增广，无监督/半监督/长尾/噪声/联邦学习)等；共发表和审稿中的15余篇SCI国际期刊和顶级会议论文，包括NeurIPS，CVPR, ECCV，ICLR，AAAI, ICASSP等CCF-A/B类会议。发明专利授权2项。

长期担任计算机视觉、人工智能、多媒体领域顶级会议CVPR, ECCV, NeurIPS, AAAI, ACM MM等审稿人。指导研究生本科生发表SCI, EI，CCF-C类会议和毕业论文累计30余篇，有丰富的保研,申博等方面经验，成功辅导学员赴南洋理工，北大，浙大等深造。

涉及范围：CCF会议A类/SCI一区、CCF会议B类/SCI二区、CCF会议C类/SCI三区、SCI四区、EI期刊、EI会议、核心期刊、研究生毕业设计

报名请扫描下方二维码了解详细情况，备注：“论文班报名”。

如果有其他想要当论文指导老师的朋友，请发简历给我，同样扫描上方二维码，备注：“论文指导老师”。基本条件：已发表两篇以上一作顶会，或3-5篇其他级别的一作论文，学历在985博士及以上。