关于yolov5模型转换为rknn

最新推荐文章于 2024-06-13 17:40:14 发布

123小星星

最新推荐文章于 2024-06-13 17:40:14 发布

阅读量3.1k

点赞数 3

文章标签： python 后端

本文链接：https://blog.csdn.net/weixin_44220689/article/details/127110943

版权

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这两个参数即可注释掉；

还有后面这部分代码也要同时注释掉；

④ 执行命令为：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周左右，感谢相互帮助的伙伴，如内容有问题还请留言多指教。