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 文件。是否需要?

相关推荐
AI浩5 小时前
PAB-Mamba-YoLo: VSSM 辅助 YOLO 用于断奶仔猪攻击行为检测
yolo
王哈哈^_^13 小时前
【完整源码+数据集】草莓数据集,yolov8草莓成熟度检测数据集 3207 张,草莓成熟度数据集,目标检测草莓识别算法系统实战教程
人工智能·算法·yolo·目标检测·计算机视觉·视觉检测·毕业设计
王哈哈^_^1 天前
YOLOv11视觉检测实战:安全距离测算全解析
人工智能·数码相机·算法·yolo·计算机视觉·目标跟踪·视觉检测
深度学习lover1 天前
<数据集>yolo航拍交通目标识别数据集<目标检测>
人工智能·python·yolo·目标检测·计算机视觉·航拍交通目标识别
Coovally AI模型快速验证1 天前
视觉语言模型(VLM)深度解析:如何用它来处理文档
人工智能·yolo·目标跟踪·语言模型·自然语言处理·开源
王哈哈^_^1 天前
【数据集+完整源码】水稻病害数据集,yolov8水稻病害检测数据集 6715 张,目标检测水稻识别算法实战训推教程
人工智能·算法·yolo·目标检测·计算机视觉·视觉检测·毕业设计
像风一样的男人@2 天前
python --两个文件夹文件名比对(yolo 图和label标注比对检查)
windows·python·yolo
AI纪元故事会3 天前
《目标检测全解析:从R-CNN到DETR,六大经典模型深度对比与实战指南》
人工智能·yolo·目标检测·r语言·cnn
Python图像识别3 天前
75_基于深度学习的咖啡叶片病害检测系统(yolo11、yolov8、yolov5+UI界面+Python项目源码+模型+标注好的数据集)
python·深度学习·yolo
Python图像识别3 天前
74_基于深度学习的垃圾桶垃圾溢出检测系统(yolo11、yolov8、yolov5+UI界面+Python项目源码+模型+标注好的数据集)
python·深度学习·yolo