YOLO模型格式转换：pt -＞ onnx -＞ rknn

导出 RKNPU 适配模型说明

Source

本仓库基于 https://github.com/ultralytics/ultralytics 仓库的 c9be1f3cce89778f79fb462797b8ca0300e3813d commit 进行修改,验证.

修改前的源码链接：

https://github.com/ultralytics/ultralytics/tree/c9be1f3cce89778f79fb462797b8ca0300e3813d

2.1 模型转换pt2onnx

如果想用最新版本的v8官方代码，可以尝试修改以下内容：
ultralytics/nn/moudles/head.py

Detect

复制代码

    def forward(self, x):
        # ============================↓新增↓============================
        if self.export and self.format == 'rknn':
            y = []
            for i in range(self.nl):
                y.append(self.cv2[i](x[i]))
                cls = torch.sigmoid(self.cv3[i](x[i]))
                cls_sum = torch.clamp(cls.sum(1, keepdim=True), 0, 1)
                y.append(cls)
                y.append(cls_sum)
            return y
        # ============================↑新增↑============================

Segment

复制代码

    def forward(self, x):
        """Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
        p = self.proto(x[0])  # mask protos
        bs = p.shape[0]  # batch size
        # ============================↓新增↓============================
        if self.export and self.format == 'rknn':
            mc = [self.cv4[i](x[i]) for i in range(self.nl)]
        else:
        # ============================↑新增↑============================        
            mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)  # mask coefficients

        x = self.detect(self, x)
        if self.training:
            return x, mc, p
        # ============================↓新增↓============================
        if self.export and self.format == 'rknn':
            bo = len(x)//3
            relocated = []
            for i in range(len(mc)):
                relocated.extend(x[i*bo:(i+1)*bo])
                relocated.extend([mc[i]])
            relocated.extend([p])
            return relocated
        # ============================↑新增↑============================               
        return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))

ultralytics\engine\exporter.py

增加rknn相关

115行

复制代码

        ['RKNN', 'rknn', '_rknnopt.torchscript', True, False],

187行

复制代码

        jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn, rknn = flags

327行

复制代码

        if rknn:
            f[12], _ = self.export_rknn()

390行左右onnx_expert附近增加函数

复制代码

    @try_export
    def export_rknn(self, prefix=colorstr('RKNN:')):
        """YOLOv8 RKNN model export."""
        LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')

        # ts = torch.jit.trace(self.model, self.im, strict=False)
        # f = str(self.file).replace(self.file.suffix, f'_rknnopt.torchscript')
        # torch.jit.save(ts, str(f))

        f = str(self.file).replace(self.file.suffix, f'.onnx')
        opset_version = self.args.opset or get_latest_opset()
        torch.onnx.export(
            self.model,
            self.im[0:1,:,:,:],
            f,
            verbose=False,
            opset_version=12,
            do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
            input_names=['images'])

        LOGGER.info(f'\n{prefix} feed {f} to RKNN-Toolkit or RKNN-Toolkit2 to generate RKNN model.\n' 
                    'Refer https://github.com/airockchip/rknn_model_zoo/tree/main/models/CV/object_detection/yolo')
        return f, None

ultralytics/nn/autobackend.py

增加rknn相关

120行左右

复制代码

        pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn, triton, rknn = \
            self._model_type(w)

550行左右

复制代码

        elif getattr(self, 'rknn', False):
            assert "for inference, please refer to https://github.com/airockchip/rknn_model_zoo/"

ultralytics/cfg/default.yaml # Export settings改为： ```format: rknn``` ultralytics/data/augment.py

复制代码

        # Create new boxes
        return np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1)), dtype=bboxes.dtype).reshape(4, n).T
        改为：
        return np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T

模型差异

在基于不影响输出结果, 不需要重新训练模型的条件下, 有以下改动:

修改输出结构, 移除后处理结构. (后处理结果对于量化不友好)
dfl 结构在 NPU 处理上性能不佳，移至模型外部的后处理阶段，此操作大部分情况下可提升推理性能。
模型输出分支新增置信度的总和，用于后处理阶段加速阈值筛选。

以上移除的操作, 均需要在外部使用CPU进行相应的处理. (对应的后处理代码可以在 RKNN_Model_Zoo 中找到)

导出onnx模型

运行pt2onnx.py

2.2 转换模型onnx2rknn

所有要用到的压缩包都放在2.AI模型转换\ubuntu中了

下载模型转换工具
网盘下载链接：https://pan.baidu.com/s/1_PquxW2rFuf77q6mT3gkDQ 提取码：6666
文件：rknn-toolkit-1.7.3-docker.tar.gz
把压缩包移到电脑端ubuntu20.04的rknn-toolkit目录，无需解压
在该目录打开终端
执行以下指令加载模型转换工具docker镜像:
复制代码
```
sudo docker load --input rknn-toolkit-1.7.3-docker.tar.gz
```
执行以下指令进入镜像bash环境测试doker是否成功加载（可跳过）:
复制代码
```
sudo docker run -t -i --privileged rknn-toolkit:1.7.3 /bin/bash
```
输入python加载python相关库
输入import rknn导入rknn相关库
没报错，表示模型转换工具环境搭建完成。

上面3步仅第一次需要配置

重新打开终端

启动doker容器，环境为rknn-toolkit:1.7.3，同时把本地model_convert文件夹映射到doker的test文件夹

这里一定要用绝对路径，Ctrl+L查看绝对路径

复制代码

sudo docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb -v /home/developer/yolov8_model_convert:/test rknn-toolkit:1.7.3 /bin/bash

进入test文件夹
quant_dataset文件夹放在test文件夹下

生成量化图片路径列表pic_path.txt，如果打开检查路径对，就不用跑这行代码，可以删掉一些，留10来行就行

复制代码

python gen_list.py

把onnx文件放目录下转rknn，注意用v8官方代码转成onnx时opset=9

复制代码

python rknn_convert.py