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 protosbs = 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 coefficientsx = self.detect(self, x)if self.training:return x, mc, p# ============================↓新增↓============================if self.export and self.format == 'rknn':bo = len(x)//3relocated = []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_exportdef 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=Falseinput_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 boxesreturn 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步仅第一次需要配置
4. 重新打开终端
启动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