当前位置：首页 > news >正文

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

news 2025/10/17 5:48:59

根据说明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 推理格式的模型。

✅ 前提条件

确保已安装必要依赖：

pip3 install -U ultralytics pnnx ncnn

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

🔧 步骤详解与脚本

步骤 1：导出原始 TorchScript 模型

yolo export model=yolov8n.pt format=torchscript

输出：yolov8n.torchscript

步骤 2：用 PNNX 初步转换（用于生成 Python 模型代码）

pnnx yolov8n.torchscript

此步骤会生成：

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

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

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

打开 yolov8n_pnnx.py，定位到类似以下代码段（通常在 forward 函数末尾）：

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 或后处理

替换为（支持动态尺寸 + 仅输出原始检测头）：

# 动态 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 末尾添加以下代码：

def export_torchscript():import torchmodel = 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，则改为：

def export_torchscript():import torchmodel = 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

python3 -c "import yolov8n_pnnx; yolov8n_pnnx.export_torchscript()"

输出：yolov8n_pnnx.py.pt

步骤 6：用 PNNX 转换为 NCNN（支持动态尺寸）

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：重命名模型文件

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