【边缘计算】YOLO模型量化及其在RK3588边缘盒子（ARM64）上的部署

西瓜和拾月

已于 2024-10-09 09:34:55 修改

阅读量4.6k

点赞数 36

文章标签：边缘计算 YOLO 人工智能 python

于 2024-09-25 14:47:44 首次发布

本文链接：https://blog.csdn.net/csdnn345/article/details/142479216

版权

一、背景

经过半年多的努力，基于GPU主机的智能检测项目已基本成型，通过在现场测试的效果来看，也基本满足用户需求，虽然其中一些功能细节仍需优化和长期测试，但下一步的工作安排直接转向了边端设备，没办法只能暂时搁置GPU主机的工作，开始搞边缘计算。出于实际性能要求（使用时对于实时性要求较高）和通用性考虑，最终选定搭载6T算力的RK3588盒子作为GPU主机的替代设备，由于自己也是初次接触边缘盒子，加之身边也没有人能指导工作，于是开始了漫长的摸索之路。

二、具体工作

0. 工具

RKNN-ToolKit2：https://github.com/airockchip/rknn-toolkit2/tree/master

1. 思路

首先需要明确工作思路，RK3588 NPU上支持的模型格式为.rknn，在模型推理时可以借助NPU进行加速，GPU主机上训练YOLO生成的权重文件为.pt格式，作为base model，后续工作都要基于这个base model展开。pt文件转rknn之前需要先转为一个通用中间格式onnx，YOLO官方代码库提供了转换脚本，只需简单修改即可完成转换。

拿到onnx文件之后，就需要使用rknn-toolkit2将onnx转为rknn，转换时根据精度和速度要求可以选择是否进行模型量化，转换完成之后就可以在盒子上进行推理精度测试和后续的业务逻辑了，盒子上需要安装的工具包是rknn-toolkit-lite2。

再次明确一下工作思路，pt-->onnx-->rknn(quantization Y/N)-->inference-->deploy

思路明确了，接下来详细展开每一步。

2. PT—>ONNX

（1）转换工具

Yolov5、v7的代码仓库里提供了转换脚本export.py，通过指定--weights转换权重文件，默认转换为torchscript，修改--include参数为onnx即可，注意--opset参数要根据安装的rknn-toolkit2版本修改算子集版本，如果rknn-toolkit2版本低，而opset版本高，会出现不匹配问题。

（2）代码修改

转换之前，我们需要拿到的是不同尺度分支送入Dectec层之前的最后一个Conv层的输出，即不同网格下的原始输出，形状为 $(bs,n_{anchors}*(n_{class}+5(x y w h conf)),ny,nx)$ ，后处理部分在NPU上进行。

修改models/yolo.py的Detect()，返回最后一个Conv层的输出

class Detect(nn.Module):

    def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
        super(Detect, self).__init__()
        self.nc = nc  # number of classes
        self.no = nc + 5  # number of outputs per anchor
        self.nl = len(anchors)  # number of detection layers
        self.na = len(anchors[0]) // 2  # number of anchors
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv

    def forward(self, x):
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv

        return x

对于yolov8，其代码结构与v5v7差别较大，不过关键还是要找到Detect层把输出截断，瑞芯微也提供了官方代码，找到Detect()，替换原v8的代码即可。实测修改之后可以成功转换rknn，只是简单测试了模型转换，没有后续使用。

（3）结构查看

修改完之后运行export.py可以导出onnx文件，使用netron查看模型结构

3. ONNX—>RKNN

（1）转换工具

PC上安装rknn-toolkit2，根据自己的python环境选择对应版本

安装完成后导入即可，如果报错可能是版本不匹配，要重新安装。

from rknn.api import RKNN

（2）转换流程

（3）转换代码

从创建rknn对象到导出rknn模型

ONNX_MODEL = 'yolov5s_relu.onnx'
RKNN_MODEL = 'yolov5s_relu.rknn'
IMG_PATH = './bus.jpg'
DATASET = './dataset.txt'

QUANTIZE_ON = True

# Create RKNN object
rknn = RKNN(verbose=True)

# pre-process config
print('--> Config model')
# 模型配置接口，确保模型转换的正确性和性能，包括设置输入均值、归一化值、是否量化，目标平台等等
rknn.config(mean_values=[[0, 0, 0]], 
            std_values=[[255, 255, 255]], 
            target_platform='rk3588')
print('done')

# Load ONNX model
print('--> Loading model')
# 加载onnx文件
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')

# Build model
print('--> Building model')
# 构建rknn模型，可以选择是否量化
ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
if ret != 0:
    print('Build model failed!')
    exit(ret)
print('done')

# Export RKNN model
print('--> Export rknn model')
# 导出rknn模型，设置导出路径
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    print('Export rknn model failed!')
    exit(ret)
print('done')

rknn.relase()

4. Quantization

rknn-toolkit2提供两种量化方式（Per-Layer 量化和 Per-Channel 量化）和三种量化算法（Normal，KL-Divergence和 MMSE）。

（1）两种量化方式

（2）三种量化算法

对精度损失依次减小，量化过程逐渐加长

（3）量化校正集

量化之前需要准备量化校正集，用于计算激活值的量化范围，防止量化跑偏。官方建议的校正集规模在20-200张，对于前两种量化算法，准备100张左右的图像基本够用，如果使用MMSE，准备30-100张即可，实测100+跟100左右的量化精度基本无差别，但再增加图像反而会造成精度下降。校正集的图像随机选取训练集或验证集中的图像即可，如果识别场景多样，则应保证选取的图像尽量覆盖模型实际应用场景。

（4）配置方法

模型转换代码中的config接口和build接口分别用于指定量化算法（默认为normal）和加载量化校正集。如果QUANTIZE_ON=False，则转换后的模型为fp16，精度几乎无损，但推理速度下降非常严重，如果实际应用场景对检测帧率要求不高，可以考虑不做量化直接转换。

# quantized_algorithm：量化算法
rknn.config(mean_values=[[0, 0, 0]], 
            std_values=[[255, 255, 255]], 
            quantized_algorithm='normal',
            target_platform='rk3588')

# do_quantization：量化开关，dataset：量化校正集
ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)

5. 混合量化

以上提到的量化配置方法为全量化，即将模型所有层的参数量化为INT8，如果该方法量化后的模型精度无法满足实际需求（此处也是当时花费时间最多的地方，要在满足最低检测帧率的情况下降低精度损失，谈何容易），则需要考虑混合量化（INT8+FP16），人为指定对模型量化过程中精度损失较为严重的层不做量化操作。

（1）混合量化流程

第一步：生成混合量化配置文件

第二步：修改配置文件，自定义量化层并导出

（2）混合量化配置

第一步：生成混合量化配置文件

调用混合量化接口hybrid_quantization_step1

rknn = RKNN()
if not os.path.exists(ONNX_MODEL):
    print('model not exist')
    exit(-1)

# pre-process config
print('--> Config model')
rknn.config(mean_values=[[0, 0, 0]],
            std_values=[[255, 255, 255]],
            optimization_level=2,
            quantized_algorithm='mmse',
            target_platform='rk3588',)
print('done')

# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')
# 生成配置文件
ret = rknn.hybrid_quantization_step1(dataset=DATASET)
rknn.release()

第二步：修改配置文件，自定义量化层并导出

第一步生成的文件包括.cfg配置文件、.model模型文件和.data数据文件，其中.cfg中保存了模型的所有层及量化参数信息，需要修改的也是该文件。

首先在全量化过程中调用精度分析接口查看量化后模型各个层的精度损失情况

# Accuracy analysis
print('--> Accuracy analysis')
ret = rknn.accuracy_analysis(inputs=['./dog_bike_car_300x300.jpg'], output_dir=None)
if ret != 0:
    print('Accuracy analysis failed!')
    exit(ret)
print('done')

找到损失较大的层，假设是743

则在.cfg文件中，在custom_quantize_layers下添加该层，参数类型设置为FP16

具体混多少层需要自己多测试验证，也可根据模型结构对该层的前后层做混合，因为损失是逐层积累的，理论上混合输入层能降低当前层损失，混合输出层能降低后续损失。使用netron打开.model文件，可以根据模型结构找到需要关闭量化的层名称。

.cfg文件修改完后，就可以调用hybrid_quantization_step2接口构建混合量化RKNN模型并导出。

rknn = RKNN()

# Build model
print('--> hybrid_quantization_step2')
ret = rknn.hybrid_quantization_step2(model_input='{model_name}.model',
                                     data_input='{model_name}.data',
                                     model_quantization_cfg='{model_name}.quantization.cfg')
if ret != 0:
    print('hybrid_quantization_step2 failed!')
    exit(ret)
print('done')

# Export RKNN model
print('--> Export RKNN model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    print('Export rknn failed!')
    exit(ret)
print('done')

6. Inference

模型量化完后就可以放到盒子上推理了。

（1）环境安装

安装rknn-toolkit-lite2。rknn-toolkit-lite2为用户提供板端模型推理的python接口，方便用户使用python语言进行AI应用开发。安装方式跟toolkit2相同。

pip3 install rknn_toolkit_lite2-1.x.y-cp39-cp39-linux_aarch64.whl

安装RKNPU2 [1]。RKNPU2提供访问RKNPU的接口——动态链接库librknnrt.so和C头文件rknn_api.h，可供用户编写C++版本的AI应用并利用RKNPU2加速；rknn-toolkit-lite2也依赖RKNPU2的librknnrt.so文件。

下载RKNPU2后将需要的文件移动到对应文件夹中并赋予执行权限即可。

# cd path/to/rknpu2/librknn_api
# cp aarch64/* /usr/lib
# cp include/rknn_api.h /usr/include
# cd ../rknn_server
# cp * /usr/bin
# chmod +x /usr/bin/rknn_server
# chmod +x /usr/bin/start_rknn.sh
# chmod +x /usr/bin/restart_rknn.sh

导入查看是否安装成功

from rknnlite.api import RKNNLite

（2）运行测试

运行rknn-toolkit-lite2下的示例代码，如果安装没问题会正常输出以下内容

至此，从模型转换到模型量化到模型边端推理的流程已走通，其它细节上的问题（比如如何进一步提高量化精度）可参考官方用户手册并结合自己的任务场景特点做优化和调整，本人也没有完全吃透，还要继续学习。

7. Deploy

板端推理测试通过后，就可以根据自己的业务逻辑编写程序了，在提升帧率方面，本人主要参考这位大佬的文章[2]，采用线程池异步的方式提高NPU占用率，进而提高检测帧率。不过本人水平有限，目前结合业务逻辑，两个模型只能做到15fps（yolov7-pose）-20fps（yolov5s）的速度，远远达不到大佬们实现的30+甚至更高的帧率。