YOLOv5s 从 PyTorch 到 RK3588 NPU 部署全流程（附性能对比）

以yolov5s-5.0为例完整记录.pt → ONNX → RKNN的转换、代码修改及 RK3588 平台推理测试。一、环境与硬件准备组件说明训练好的模型yolov5s.ptv5.0转换设备x86_64 架构 LinuxUbuntu 20.04 Python 3.8 演示部署设备RK3588 开发板二、PT → ONNX 转换1. 修改模型forward函数适配 NPU 导出在 YOLOv5 的models/yolo.py中找到Detect类的forward方法修改如下修改前原始def forward(self, x): Processes input through YOLOv5 layers, altering shape for detection: x(bs, 3, ny, nx, 85). z [] # inference output for i in range(self.nl): x[i] self.m[i](x[i]) # conv bs, _, ny, nx x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) x[i] x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() if not self.training: # inference if self.dynamic or self.grid[i].shape[2:4] ! x[i].shape[2:4]: self.grid[i], self.anchor_grid[i] self._make_grid(nx, ny, i) if isinstance(self, Segment): # (boxes masks) xy, wh, conf, mask x[i].split((2, 2, self.nc 1, self.no - self.nc - 5), 4) xy (xy.sigmoid() * 2 self.grid[i]) * self.stride[i] # xy wh (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i] # wh y torch.cat((xy, wh, conf.sigmoid(), mask), 4) else: # Detect (boxes only) xy, wh, conf x[i].sigmoid().split((2, 2, self.nc 1), 4) xy (xy * 2 self.grid[i]) * self.stride[i] # xy wh (wh * 2) ** 2 * self.anchor_grid[i] # wh y torch.cat((xy, wh, conf), 4) z.append(y.view(bs, self.na * nx * ny, self.no)) return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)修改后增加环境变量控制def forward(self, x): z [] # inference output for i in range(self.nl): if os.getenv(RKNN_model_hack, 0) ! 0: x[i] torch.sigmoid(self.m[i](x[i])) # conv return x作用通过环境变量RKNM_model_hack控制导出时仅保留 sigmoid 激活后的卷积输出简化 ONNX 结构。2. 修改export.py脚本在export.py开头增加import os os.environ[RKNM_model_hack] npu_2修改模型输出 shape 获取方式修改前shape tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape修改后shape tuple(y[0].shape) # model output shape3. 执行 ONNX 导出python export.py --weights yolov5s.pt --include onnx得到yolov5s.onnx文件。三、ONNX → RKNN 转换在 x86 Linux 上完成⚠️ 注意此步骤不能在 RK3588 上执行需在 x86_64 主机或虚拟机/WSL中进行。1. 创建 Conda 环境Python 3.8conda create -n RKN python3.8 -y conda activate RKN2. 安装 rknn_toolkit2git clone https://github.com/rockchip-linux/rknn-toolkit2.git cd rknn-toolkit2 pip install rknn_toolkit2-1.4.0_22dcfef4-cp38-linux_x86_64.whl测试安装from rknn.api import RKNN3. 修改转换示例脚本进入示例目录cd rknn-toolkit2-1.6.0/rknn-toolkit2/examples/onnx/yolov5将你的yolov5s.onnx放入该目录修改test.py中的模型路径和图片路径执行转换python3 test.py成功后会生成.rknn模型文件。四、在 RK3588 平台上部署并推理1. 下载 rknpu2 工具git clone https://github.com/rockchip-linux/rknpu2 cd rknpu2-master/examples/rknn_yolov5_demo2. 修改头文件设定类别数进⼊到 /rknpu2-master/examples/rknn_yolov5_demo⽂件夹下修改头⽂件postprocess.h#define OBJ_CLASS_NUM 2 // 修改为你的数据集类别数3. 修改标签文件替换model/coco_80_labels_list.txt内容为你的类别名称。4. 放置模型和测试图片将.rknn文件放到model/RK3588/目录下将测试图片如test.jpg放到model/目录下5. 编译并运行./build-linux_RK3588.sh # 生成 install 目录 cd install/rknn_yolov5_demo_linux ./rknn_yolov5_demo ./model/RK3588/best.rknn ./model/test.jpg五、推理性能对比RK3588 实测测试条件同一张测试图片循环推理 10 次取平均计时范围预处理 推理 后处理NPU 单核占用约30% ~ 44%模型模型大小平均延时yolov5s-cpuCPU 推理28.5 MB398 ~ 456 msyolov5s-npuNPU 推理8.2 MB27 ~ 29 msyolov8sCPU 推理12.2 MB224 ~ 240 ms NPU 加速效果显著延时从数百毫秒降至30ms 以内模型体积也大幅缩小。