关于yolov5模型转换为rknn

1.首先是版本问题
由于硬件为RK1808,本次采用rknn-toolkit 1.7.1进行模型转换和验证。
将pt模型部署在硬件上需要使用C++编程部署,并且需要将pt模型转化为RK可接受的模型。
① 目前硬件环境RK仅能解析yolov5 5.*版本的模型,对于yolov5更新后的6.0版本,硬件并不支持。
因为yolov5 5.*模型输出形式为三组数据,而yolov5 6.*输出为一组数据。(个人觉得是后期数据处理不同,并不影响检测结果的输出,但是要用于rknn的部署,就麻烦很多)
② 关于rknn转换代码,目前采用rknn_model_zoo能够实现对yolov5 模型的转化,相较之前转换文档,该文档参数清晰明了,方便操作。
2.转换流程
1)训练pt模型best.pt,本次采用yolov5 5.*版进行训练及转换,通过yolov5官方文档中的./model/export.py转化为onnx模型,即best.onnx。
如果在转换过程中出现下面的报错:

AttributeError: Can‘t get attribute ‘SPPF‘ on <module ‘models.common‘ from ‘model\\common.py‘

则需要修改common.py文件,在里面参考博文https://blog.csdn.net/magicnnonn/article/details/124343085添加如下代码:

import warnings
class SPPF(nn.Module):
    # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
    def __init__(self, c1, c2, k=5):
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
 
    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            y1 = self.m(x)
            y2 = self.m(y1)
            return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))

2) onnx模型验证,该过程需要在rknn-toolkit 1.7.1环境下进行。进入rknn_model_zoo/models/vision/object_detection/yolov5-pytorch/RKNN_python_demo,文件yolo_map_test_rknn.py可以验证best.onnx模型的预测情况。
若训练自己的数据集和类别,需要将yolo_map_test_rknn.py中的部分内容进行修改:
① 下面三个参数需要根据需要进行修改,OBJ_THRESH为置信度阈值,NMS_THRESH为NMS阈值,IMG_SIZE为模型输入图像大小
在这里插入图片描述
② CLASSES为类别名称,可根据实际换成自己预测的类别名称。
在这里插入图片描述
③ 若不采用coco_map进行模型验证,coco_id_list和CLASSESCOCO这两个参数即可注释掉;
关于yolov5模型转换为rknn_第1张图片
还有后面这部分代码也要同时注释掉;
关于yolov5模型转换为rknn_第2张图片
④ 执行命令为:python yolo_map_test_rknn.py --model_path ./best_test.onnx --img_show --img_folder ./img
其中: model_path 为待验证模型路径,img_show展示预测结果,img_folder为待预测图像路径
该预测结果可于pt模型的预测结果进行对比。(在个别图片上检测结果略有偏差,该偏差为中性,检测结果变好变差都有可能。这是因为yolov5.detect.py的实现允许动态尺寸输入导致的,RKNN模型目前不允许动态尺寸输入,所以最终的检测结果存在部分差异。)
3)onnx模型转rknn模型;
模型转换程序所在位置为./rknn_model_zoo/common/rknn_converter/
①在这里有一个小小的改动,将config_parser.py 中的batch_size设置为了1
在这里插入图片描述
②可以参照./rknn_model_zoo/common/rknn_converter/README.md中的各个参数特性,创建符合自己模型情况和数据的yml参数文件
这里需要特别注意mean_values和 std_values两个参数的设置,
如果需要经输入数据归一化在[0,1]之间设置为mean_values:0,0,0;std_values:255 ;
如果需要将输入数据归一化到[-1, 1]之间,则可以设置为mean_values:128,128,128;std_values:128 。
③运行命令:python ./common/rknn_convert.py --yml_path ./common/data/onnx_model_config.yml --output_path ./rknn_models/best_test.rknn --compute_convert_loss --eval_perf
4)rknn模型验证与onnx模型验证一致,如果onnx验证时已经改过参数则不需要再修改, 执行命令为:python yolo_map_test_rknn.py --model_path ./best_test.rknn --img_show --img_folder ./img_rk
本次笔记相对完整,耗时2周左右,感谢相互帮助的伙伴,如内容有问题还请留言多指教。

你可能感兴趣的:(python,后端)