pnnx

PyTorch 神经网络交换格式

注：当前实现位于 https://github.com/Tencent/ncnn/tree/master/tools/pnnx

什么是 pnnx

PyTorch 神经网络交换格式（PNNX）是用于 PyTorch 模型互操作性的开放标准。PNNX 为 PyTorch 提供了一种开放的模型格式，它严格遵循 PyTorch 的定义来描述计算图和高级算子。

优化 Torch 模型	减少扩展包依赖
将 TorchScript / ONNX 转换回 Python	导出为便携式 PNNX 格式 / ONNX-Zero / NCNN

flowchart TD
    torchmodel["torch 模型<br>torchvision.models.resnet18()"]
    othermodel["caffe, mxnet<br>keras, tensorflow<br>paddlepaddle, 等"]
    torchscript["torchscript 文件<br>resnet18.pt"]
    onnx["onnx 文件<br>resnet18.onnx"]
    tnn["tnn 文件<br>resnet18.tnnproto/tnnmodel"]
    optmodel["优化后的 torch 模型<br>resnet18_pnnx.Model()"]
    ncnnmodel["ncnn 模型<br>resnet18.ncnn.param/bin"]
    onnxzeromodel["onnx-zero 模型<br>resnet18.pnnx.onnx"]

    subgraph pnnx
        optmodel
        ncnnmodel
        onnxzeromodel
    end

    torchmodel -->|"mod = torch.jit.trace(model, x)<br>mod.save('resnet18.pt')"| torchscript
    torchmodel -->|"torch.onnx.export(model, x, 'resnet18.onnx')"| onnx
    othermodel -->|"导出到 onnx"| onnx
    torchmodel -->|"pnnx.export(model, 'resnet18.pt', x)"| pnnx
    torchscript -->|"pnnx.convert('resnet18.pt', x)"| pnnx
    onnx -->|"pnnx.convert('resnet18.onnx', x)"| pnnx
    tnn -->|"pnnx.convert('resnet18.tnnproto', x)"| pnnx

如何安装 pnnx

----- A. python pip（推荐）

• Windows/Linux/macOS 64位
• python 3.7 或更高版本

pip3 install pnnx

----- B. 便携式二进制包（如果你不喜欢 Python，推荐使用）

• Windows/Linux/macOS 64位
• 对于 Linux，glibc 2.17+

从 https://github.com/pnnx/pnnx/releases 下载便携式 pnnx 二进制包并解压。

该包包含了所有必要的二进制文件。它是便携式的，因此不需要 CUDA 或 PyTorch 运行时环境 :)

Windows

Linux(x86_64)

Linux(aarch64)

macOS(universal)

----- C. 从源码构建

1. 安装 PyTorch
2. （可选）安装 torchvision 以支持 pnnx 的 torchvision 算子
3. （可选）安装 protobuf 以支持 pnnx 的 onnx-zero
4. 克隆 https://github.com/Tencent/ncnn.git
5. 使用 cmake 在 ncnn/tools/pnnx 中构建 pnnx

你可能会参考 https://github.com/pnnx/pnnx/blob/main/.github/workflows/release.yml 获取详细步骤

git clone https://github.com/Tencent/ncnn.git
mkdir ncnn/tools/pnnx/build
cd ncnn/tools/pnnx/build
cmake -DCMAKE_INSTALL_PREFIX=install -DTorch_INSTALL_DIR=<your libtorch install dir> -DTorchVision_INSTALL_DIR=<your torchvision install dir> ..
cmake --build . --config Release -j 4
cmake --build . --config Release --target install

如何使用 pnnx

----- A. python

1. 使用 pnnx.export() 优化并导出你的 torch 模型

import torch
import torchvision.models as models
import pnnx

model = models.resnet18(pretrained=True)

x = torch.rand(1, 3, 224, 224)

opt_model = pnnx.export(model, "resnet18.pt", x)

# 对于多输入模型，可以使用元组
# opt_model = pnnx.export(model, "resnet18.pt", (x, y, z))

2. 像使用普通模块一样使用优化后的模块

result = opt_model(x)

3. 选择 resnet18_pnnx.py 作为 pnnx 优化后的 torch 模型
4. 选择 resnet18.ncnn.param 和 resnet18.ncnn.bin 用于 ncnn 推理

----- B. 命令行

1. 将你的 torch 模型导出为 torchscript / onnx

import torch
import torchvision.models as models

net = models.resnet18(pretrained=True)
net = net.eval()

x = torch.rand(1, 3, 224, 224)

# 如果在跟踪时出现错误，可以尝试禁用检查
# mod = torch.jit.trace(net, x, check_trace=False)
mod = torch.jit.trace(net, x)

mod.save("resnet18.pt")

# 也可以尝试导出为传统的 onnx 格式
torch.onnx.export(net, x, 'resnet18.onnx')

2. 使用 pnnx 将 torchscript / onnx 转换为优化后的 pnnx 模型和 ncnn 模型文件

./pnnx resnet18.pt inputshape=[1,3,224,224]
./pnnx resnet18.onnx inputshape=[1,3,224,224]
./pnnx resnet18.tnnproto inputshape=[1,3,224,224]

macOS zsh 用户可能需要使用双引号以避免歧义

./pnnx resnet18.pt "inputshape=[1,3,224,224]"

对于多输入模型，使用列表

./pnnx resnet18.pt inputshape=[1,3,224,224],[1,32]

对于非 fp32 输入数据类型的模型，添加类型后缀

./pnnx resnet18.pt inputshape=[1,3,224,224]f32,[1,32]i64

3. 选择 resnet18_pnnx.py 作为 pnnx 优化后的 torch 模型
4. 选择 resnet18.ncnn.param 和 resnet18.ncnn.bin 用于 ncnn 推理

----- 使用 netron 可视化 pnnx

在浏览器中打开 https://netron.app/，然后将 resnet18.pnnx.param 或 resnet18.pnnx.onnx 拖入其中。

----- 比较推理结果

通常情况下，你会得到七个文件：

resnet18.pnnx.param	PNNX 图定义
resnet18.pnnx.bin	PNNX 模型权重
resnet18_pnnx.py	用于推理的 PyTorch 脚本，包含模型构建和权重初始化的 Python 代码
resnet18.pnnx.onnx	以 ONNX 格式表示的 PNNX 模型
resnet18.ncnn.param	NCNN 图定义
resnet18.ncnn.bin	NCNN 模型权重
resnet18_ncnn.py	用于推理的 pyncnn 脚本

pip3 install ncnn

运行由 PNNX 生成的推理脚本，并检查输出是否足够接近。

$ python resnet18_pnnx.py
tensor([[-1.5752e+00,  6.8381e-01,  1.4599e+00,  1.1986e+00,  1.0503e+00,
         -5.0585e-01,  9.1962e-01,  1.4127e-01, -1.2256e+00, -3.8200e-01,
          1.1840e+00,  2.5817e+00,  1.3319e+00,  2.8250e+00,  2.6328e+00,
          1.1827e+00,  1.6862e+00,  2.9742e-01,  1.5851e+00,  1.9562e+00,
          ...

$ python resnet18_ncnn.py
tensor([[-1.5719e+00,  6.8591e-01,  1.4592e+00,  1.1973e+00,  1.0503e+00,
         -5.0833e-01,  9.1693e-01,  1.4180e-01, -1.2239e+00, -3.8417e-01,
          1.1816e+00,  2.5768e+00,  1.3295e+00,  2.8196e+00,  2.6259e+00,
          1.1806e+00,  1.6830e+00,  2.9536e-01,  1.5808e+00,  1.9530e+00,
          ...

----- 详细选项

用法：pnnx [model.pt] [(key=value)...]
  pnnxparam=model.pnnx.param
  pnnxbin=model.pnnx.bin
  pnnxpy=model_pnnx.py
  pnnxonnx=model.pnnx.onnx
  ncnnparam=model.ncnn.param
  ncnnbin=model.ncnn.bin
  ncnnpy=model_ncnn.py
  fp16=1
  optlevel=2
  device=cpu/gpu
  inputshape=[1,3,224,224],...
  inputshape2=[1,3,320,320],...
  customop=/home/nihui/.cache/torch_extensions/fused/fused.so,...
  moduleop=models.common.Focus,models.yolo.Detect,...
示例用法：pnnx mobilenet_v2.pt inputshape=[1,3,224,224]
              pnnx yolov5s.pt inputshape=[1,3,640,640]f32 inputshape2=[1,3,320,320]f32 device=gpu moduleop=models.common.Focus,models.yolo.Detect

参数	默认值	说明
model.pt model.onnx model.tnnproto	(必填)	TorchScript / ONNX / TNN 文件路径
pnnxparam	*.pnnx.param (* 是模型名称)	PNNX 图定义文件
pnnxbin	*.pnnx.bin	PNNX 模型权重
pnnxpy	*_pnnx.py	用于推理的 PyTorch 脚本，包含模型构建和权重初始化代码
pnnxonnx	*.pnnx.onnx	以 ONNX 格式表示的 PNNX 模型
ncnnparam	*.ncnn.param	NCNN 图定义
ncnnbin	*.ncnn.bin	NCNN 模型权重
ncnnpy	*_ncnn.py	用于推理的 pyncnn 脚本
fp16	1	以 FP16 数据类型保存 NCNN 权重和 ONNX 模型
optlevel	2	图优化级别 0 = 不应用优化 1 = 针对推理进行优化 2 = 更进一步针对推理优化
device	cpu	TorchScript 模型输入的设备类型，ONNX 模型忽略此参数，可为 cpu 或 gpu
inputshape	(可选)	模型输入的形状。用于解析模型图中的张量形状。例如，对于只有一个输入的模型，使用 [1,3,224,224]；对于有两个输入的模型，使用 [1,3,224,224],[1,3,224,224]。 形状元组可以附加数据类型后缀，如 [1,3,224,224]f32 表示 float32 类型。 f32 = torch.float32 或 torch.float f64 = torch.float64 或 torch.double f16 = torch.float16 或 torch.half u8 = torch.uint8 i8 = torch.int8 i16 = torch.int16 或 torch.short i32 = torch.int32 或 torch.int i64 = torch.int64 或 torch.long c32 = torch.complex32 c64 = torch.complex64 c128 = torch.complex128
inputshape2	(可选)	备选模型输入的形状，格式与 inputshape 相同。通常与 inputshape 一起使用，以解析模型图中的动态形状（-1）。
customop	(可选)	自定义算子的 Torch 扩展库（动态链接库）列表，ONNX 模型忽略此参数，各库之间用逗号分隔。例如，/home/nihui/.cache/torch_extensions/fused/fused.so,...
moduleop	(可选)	需要保持为一个整体算子的模块列表，ONNX 模型忽略此参数，各模块之间用逗号分隔。例如，models.common.Focus,models.yolo.Detect

PNNX 快速上手指南

PNNX (PyTorch Neural Network eXchange) 是一个专为 PyTorch 模型设计的开放标准格式。它能将 PyTorch 模型优化并转换为便携式格式，同时支持导出为 ncnn（高性能推理框架）和 ONNX-Zero 格式，旨在解决 PyTorch 模型部署中的依赖复杂和算子不匹配问题。

1. 环境准备

系统要求

• 操作系统: Windows / Linux / macOS (64 位)
• Python 版本: 3.7 或更高 (若使用 Python 方式)
• Linux 额外要求: glibc 2.17+ (若使用二进制包)

前置依赖

• 方式 A (Python): 需预先安装 torch 和 torchvision。
• 方式 B (二进制包): 无需安装 CUDA、PyTorch 或任何 Python 环境，开箱即用。
• 方式 C (源码编译): 需安装 CMake、PyTorch (libtorch)、可选安装 torchvision 和 protobuf。

2. 安装步骤

推荐根据实际需求选择以下任一方式：

方式 A：通过 pip 安装 (适合 Python 开发者)

pip3 install pnnx

注：国内用户如遇下载慢，可指定清华源：pip3 install pnnx -i https://pypi.tuna.tsinghua.edu.cn/simple

方式 B：下载便携二进制包 (适合无 Python 环境或追求轻量)

访问 GitHub Releases 页面下载对应系统的压缩包并解压：

• Windows: 下载链接
• Linux (x86_64): 下载链接
• Linux (aarch64): 下载链接
• macOS: 下载链接

解压后，目录内即包含可执行的 pnnx 程序。

方式 C：源码编译 (适合高级用户)

git clone https://github.com/Tencent/ncnn.git
mkdir ncnn/tools/pnnx/build
cd ncnn/tools/pnnx/build
# 请替换 <your libtorch install dir> 为你的实际路径
cmake -DCMAKE_INSTALL_PREFIX=install -DTorch_INSTALL_DIR=<your libtorch install dir> -DTorchVision_INSTALL_DIR=<your torchvision install dir> ..
cmake --build . --config Release -j 4
cmake --build . --config Release --target install

3. 基本使用

PNNX 支持两种主要工作流：直接在 Python 中调用或通过命令行转换已有模型。

场景一：在 Python 代码中直接优化并导出

适用于你拥有原始的 PyTorch 模型代码。

import torch
import torchvision.models as models
import pnnx

# 1. 加载模型
model = models.resnet18(pretrained=True)

# 2. 构造示例输入 (用于追踪计算图)
x = torch.rand(1, 3, 224, 224)

# 3. 导出模型
# 生成 resnet18.pt (PNNX 格式), resnet18.ncnn.param/bin, resnet18_pnnx.py 等文件
opt_model = pnnx.export(model, "resnet18.pt", x)

# 4. 像普通模型一样使用优化后的模块
result = opt_model(x)

产出文件说明：

• resnet18_pnnx.py: 可直接运行的 PyTorch 推理脚本。
• resnet18.ncnn.param / .bin: 用于 ncnn 推理的模型文件。

场景二：通过命令行转换现有模型

适用于你已经有导出的 TorchScript (.pt) 或 ONNX (.onnx) 模型。

第一步：准备模型文件 先使用 PyTorch 原生接口导出模型（如果尚未导出）：

import torch
import torchvision.models as models

net = models.resnet18(pretrained=True)
net = net.eval()
x = torch.rand(1, 3, 224, 224)

# 导出为 TorchScript
mod = torch.jit.trace(net, x)
mod.save("resnet18.pt")

# 或者导出为 ONNX
# torch.onnx.export(net, x, 'resnet18.onnx')

第二步：运行 PNNX 转换命令 在终端执行（假设已安装或在二进制包目录下）：

# 转换 TorchScript 模型
./pnnx resnet18.pt inputshape=[1,3,224,224]

# 转换 ONNX 模型
./pnnx resnet18.onnx inputshape=[1,3,224,224]

注意：macOS zsh 用户可能需要给参数加引号以防歧义：

./pnnx resnet18.pt "inputshape=[1,3,224,224]"

第三步：验证结果 转换完成后会生成多个文件，运行生成的 Python 脚本对比输出：

# 运行 PNNX 优化的 PyTorch 脚本
python resnet18_pnnx.py

# 运行 ncnn 推理脚本 (需先安装 ncnn: pip3 install ncnn)
python resnet18_ncnn.py

检查两者输出的 Tensor 数值是否接近，以确认转换成功。

进阶提示

• 多输入模型: 使用逗号分隔形状，例如 inputshape=[1,3,224,224],[1,32]。
• 指定数据类型: 在形状后添加后缀，如 f32 (float32), i64 (int64)，例如 inputshape=[1,3,224,224]f32。
• 可视化: 将生成的 resnet18.pnnx.param 或 resnet18.pnnx.onnx 拖入 Netron 即可查看模型结构。