PNNX + TorchScript + 手动修改后处理逻辑,最终输出适配 NCNN官方 yolov8.cpp

根据说明https://github.com/Tencent/ncnn/blob/master/examples/yolov8.cpp

请根据以下说明,给出yolov8n.pt模型转换为ncnn模型,详细步骤与脚本代码:

// Copyright 2024 Tencent

// SPDX-License-Identifier: BSD-3-Clause

// 1. install

// pip3 install -U ultralytics pnnx ncnn

// 2. export yolov8 torchscript

// yolo export model=yolov8n.pt format=torchscript

// 3. convert torchscript with static shape

// pnnx yolov8n.torchscript

// 4. modify yolov8n_pnnx.py for dynamic shape inference

// A. modify reshape to support dynamic image sizes

// B. permute tensor before concat and adjust concat axis

// C. drop post-process part

// before:

// v_165 = v_142.view(1, 144, 6400)

// v_166 = v_153.view(1, 144, 1600)

// v_167 = v_164.view(1, 144, 400)

// v_168 = torch.cat((v_165, v_166, v_167), dim=2)

// ...

// after:

// v_165 = v_142.view(1, 144, -1).transpose(1, 2)

// v_166 = v_153.view(1, 144, -1).transpose(1, 2)

// v_167 = v_164.view(1, 144, -1).transpose(1, 2)

// v_168 = torch.cat((v_165, v_166, v_167), dim=1)

// return v_168

// 5. re-export yolov8 torchscript

// python3 -c 'import yolov8n_pnnx; yolov8n_pnnx.export_torchscript()'

// 6. convert new torchscript with dynamic shape

// pnnx yolov8n_pnnx.py.pt inputshape=[1,3,640,640] inputshape2=[1,3,320,320]

// 7. now you get ncnn model files

// mv yolov8n_pnnx.py.ncnn.param yolov8n.ncnn.param

// mv yolov8n_pnnx.py.ncnn.bin yolov8n.ncnn.bin

// the out blob would be a 2-dim tensor with w=144 h=8400

//

// | bbox-reg 16 x 4 | per-class scores(80) |

// ±----±----±----±----±---------------------+

// | dx0 | dy0 | dx1 | dy1 |0.1 0.0 0.0 0.5 ...|

// all /| | | | | . |

// boxes | ... | ... | ... | ... |0.0 0.9 0.0 0.0 ...|

// (8400)| | | | | . |

// | | | | | . |

// ±----±----±----±----±---------------------+

//

根据你提供的详细说明,以下是将 yolov8n.pt 模型转换为 支持动态输入尺寸NCNN 模型 的完整步骤与脚本代码。整个流程基于 PNNX + TorchScript + 手动修改后处理逻辑 ,最终输出适配 NCNN 官方 yolov8.cpp 推理格式的模型。


✅ 前提条件

确保已安装必要依赖:

bash 复制代码
pip3 install -U ultralytics pnnx ncnn

注意:pnnx 包仅提供 Python 接口,实际转换依赖系统 PATH 中的 pnnx 可执行文件(需从 https://github.com/pnnx/pnnx 编译或下载预编译版)。


🔧 步骤详解与脚本

步骤 1:导出原始 TorchScript 模型

bash 复制代码
yolo export model=yolov8n.pt format=torchscript

输出:yolov8n.torchscript


步骤 2:用 PNNX 初步转换(用于生成 Python 模型代码)

bash 复制代码
pnnx yolov8n.torchscript

此步骤会生成:

  • yolov8n_pnnx.py
  • yolov8n_pnnx.bin
  • 其他中间文件

我们只关心 yolov8n_pnnx.py,用于手动修改。


步骤 3:修改 yolov8n_pnnx.py 支持动态尺寸并移除后处理

打开 yolov8n_pnnx.py,定位到类似以下代码段(通常在 forward 函数末尾):

python 复制代码
v_165 = v_142.view(1, 144, 6400)
v_166 = v_153.view(1, 144, 1600)
v_167 = v_164.view(1, 144, 400)
v_168 = torch.cat((v_165, v_166, v_167), dim=2)
...
return v_final  # 包含 NMS 或后处理

替换为(支持动态尺寸 + 仅输出原始检测头):

python 复制代码
# 动态 reshape + transpose
v_165 = v_142.view(1, 144, -1).transpose(1, 2)  # [1, 6400, 144]
v_166 = v_153.view(1, 144, -1).transpose(1, 2)  # [1, 1600, 144]
v_167 = v_164.view(1, 144, -1).transpose(1, 2)  # [1,  400, 144]

# 沿检测框数量维度拼接(dim=1)
v_168 = torch.cat((v_165, v_166, v_167), dim=1)  # [1, 8400, 144]

# 直接返回原始输出,不进行 sigmoid/NMS
return v_168

✅ 这样输出张量形状为 [1, 8400, 144],对应 NCNN 中的 w=144, h=8400(NCNN 是 HWC 布局,但 blob 维度为 w=channel)。


步骤 4:添加导出函数(便于重新生成 TorchScript)

yolov8n_pnnx.py 末尾添加以下代码:

python 复制代码
def export_torchscript():
    import torch
    model = Model()  # 确保类名为 Model,若不是请替换
    model.eval()
    dummy_input = torch.randn(1, 3, 640, 640)
    traced = torch.jit.trace(model, dummy_input)
    traced.save("yolov8n_pnnx.py.pt")
    print("✅ TorchScript saved as yolov8n_pnnx.py.pt")

⚠️ 如果你的模型类名不是 Model,请根据 yolov8n_pnnx.py 中的实际类名修改(如 class yolov8n_pnnx(torch.nn.Module),则应实例化 yolov8n_pnnx())。

例如,若类名为 yolov8n_pnnx,则改为:

python 复制代码
def export_torchscript():
    import torch
    model = yolov8n_pnnx()
    model.eval()
    dummy_input = torch.randn(1, 3, 640, 640)
    traced = torch.jit.trace(model, dummy_input)
    traced.save("yolov8n_pnnx.py.pt")
    print("✅ TorchScript saved as yolov8n_pnnx.py.pt")

步骤 5:重新导出修改后的 TorchScript

bash 复制代码
python3 -c "import yolov8n_pnnx; yolov8n_pnnx.export_torchscript()"

输出:yolov8n_pnnx.py.pt


步骤 6:用 PNNX 转换为 NCNN(支持动态尺寸)

bash 复制代码
pnnx yolov8n_pnnx.py.pt inputshape=[1,3,640,640] inputshape2=[1,3,320,320]

inputshape2 用于指定第二个测试尺寸,使 PNNX 生成动态兼容的 NCNN 模型。
输出:

  • yolov8n_pnnx.py.ncnn.param
  • yolov8n_pnnx.py.ncnn.bin

步骤 7:重命名模型文件

bash 复制代码
mv yolov8n_pnnx.py.ncnn.param yolov8n.ncnn.param
mv yolov8n_pnnx.py.ncnn.bin yolov8n.ncnn.bin

✅ 现在你已获得可用于 ncnn/examples/yolov8.cpp 的模型!


📌 最终输出说明

  • 模型输出 blob :形状为 (w=144, h=8400)
    • 前 4 列:dx, dy, dw, dh(需转换为 xywh 或 xyxy)
    • 后 80 列:未经过 sigmoid 的类别 logits
  • 后处理需在 C++ 中完成
    • 对 80 类别分值应用 sigmoid
    • 计算置信度 = obj_conf(隐含在 bbox 回归中) × cls_conf
    • NMS

这与官方 yolov8.cpp 完全一致。


📝 补充建议

  • 若你后续需要 INT8 量化 ,请先确保 FP32 模型推理正确,再使用 ncnn2int8
  • 若 PNNX 报错找不到算子,可尝试升级 PNNX 到最新版(YOLOv8 支持在 2023 年后逐步完善)。

如需自动化脚本,我也可以为你打包成 .sh.py 文件。是否需要?

相关推荐
智驱力人工智能1 小时前
小区高空抛物AI实时预警方案 筑牢社区头顶安全的实践 高空抛物检测 高空抛物监控安装教程 高空抛物误报率优化方案 高空抛物监控案例分享
人工智能·深度学习·opencv·算法·安全·yolo·边缘计算
工程师老罗1 小时前
YOLOv1 核心结构解析
yolo
Lun3866buzha2 小时前
YOLOv10-BiFPN融合:危险物体检测与识别的革新方案,从模型架构到实战部署全解析
yolo
Katecat996632 小时前
YOLOv8-MambaOut在电子元器件缺陷检测中的应用与实践_1
yolo
工程师老罗3 小时前
YOLOv1 核心知识点笔记
笔记·yolo
工程师老罗8 小时前
基于Pytorch的YOLOv1 的网络结构代码
人工智能·pytorch·yolo
学习3人组11 小时前
YOLO模型集成到Label Studio的MODEL服务
yolo
孤狼warrior12 小时前
YOLO目标检测 一千字解析yolo最初的摸样 模型下载,数据集构建及模型训练代码
人工智能·python·深度学习·算法·yolo·目标检测·目标跟踪
水中加点糖14 小时前
小白都能看懂的——车牌检测与识别(最新版YOLO26快速入门)
人工智能·yolo·目标检测·计算机视觉·ai·车牌识别·lprnet
前端摸鱼匠1 天前
YOLOv8 环境配置全攻略:Python、PyTorch 与 CUDA 的和谐共生
人工智能·pytorch·python·yolo·目标检测