【腾讯犀牛鸟2026】pnnx 支持 torch.export / PT2 模型转换

[mingshi2333]

腾讯犀牛鸟开源人才培养计划第三阶段（Shape with AI · 开源课题实战）项目分享。

这次做的是 pnnx for torch exported program，给 pnnx 增加直接读取 torch.export.save() 生成的 .pt2 文件的能力。转换入口是 pnnx model.pt2，读取模型图、参数和常量后，接入现有 PNNX passes，生成 PNNX Python 和 ncnn 模型。

我从 2026 年 8 月 23 日 开始这项 PT2 支持的实现（最早的开发提交），并于 8 月 28 日 提交上游 PR #6933，之后持续补充功能和处理 review。

目前已实现静态推理图转换、生成模型的 PT2 再导出，以及基本动态维度支持。原有算子测试也增加了 PT2 路径，和 TorchScript 使用相同的模型及数值比较。

代码还在上游 PR review 中，试用需要使用下面的分支：

• 代码：mingshi2333/ncnn · pnnx-exported-program
• 上游 PR：pnnx support torch exported program #6933
• 功能范围与测试入口：tools/pnnx/README.md

转换流程

ExportedProgram 中除了算子图，还保存了图输入的身份、张量类型与形状、参数、常量和动态维度约束。直接读取这些信息后，用户不需要先把模型另转成 TorchScript 或 ONNX。

PyTorch 模型
  ↓ torch.export.export / torch.export.save
model.pt2
  ↓ 读取 archive、JSON、图签名与权重
ExportedProgram 中间结构
  ↓ 规范化图、绑定 ATen 参数、生成 PNNX 节点
PNNX IR → 现有优化与 lowering
             ├─ PNNX Python → 执行 / 再导出 PT2
             └─ ncnn 模型   → 原生推理

.pt2 和 TorchScript 文件都可以是 ZIP 容器，不能只看到 ZIP 文件头就送进 TorchScript loader。现在根据 archive 中的标记识别 PT2，再检查布局和版本。已识别的 PT2 如果不支持，会报告具体错误。

这条路径处理 torch.export.save() 保存的程序。AOTInductor 编译后的运行包不在当前范围内。

前端如何拆分

新增代码主要在 tools/pnnx/src/，解析、算子绑定和图生成分开处理：

文件	职责
model_format.*、pt2_archive.*、storezip.*	格式识别、archive 布局和 ZIP 条目读取
json_reader.*	JSON 语法解析和错误位置
exported_program_schema.*	图签名、张量元数据、算子参数和动态维度等中间结构
exported_program_graph.*	图与参数的规范化
exported_program_tensor.*	根据 dtype、shape、stride 和 offset 还原权重
exported_program_operator.*	ATen overload、默认参数和别名/写入信息
load_exported_program.*	连接上述模块，检查表示范围并生成 PNNX IR

JSON 和 ZIP 部分没有新增第三方依赖。ZIP 读取扩展了 pnnx 已有的 StoreZipReader，使用 central directory 获取条目位置和大小，处理 local header 未填写大小的 data descriptor 情况，并检查边界与 CRC。实际读取的条目要求使用 STORE，不支持压缩数据解码。

JSON 解析后先进入有明确字段类型的中间结构。这样处理 schema 变动时，可以集中修改字段读取和版本适配，后面的图生成继续使用同一套结构。代码仍会检查版本，未知 schema 不会因为“有兼容层”就自动接受。

权重与普通输入的区别

ExportedProgram 会把参数和 buffer 提升成图输入。导入时需要结合 graph signature 判断身份：用户 Tensor 变成 pnnx.Input，参数、buffer 和 Tensor 常量则变成模型属性。

权重也不能一律当作连续数组复制。共享 storage、非零 offset、转置后的 stride 和字节序都要参与还原。PNNX 保存张量的值、形状和 dtype，但不承诺还原原始 Python Module 的 state_dict 命名、buffer 持久性或训练状态。

PNNX Python 和原生 ncnn 的精度边界分别处理。例如 Double state 在 PNNX 中保留 float64，进入原生 ncnn lowering 时转成 fp32，避免把每元素 8 字节的数据交给按 4 字节读取的原生层。

ATen 参数如何补齐

PT2 中的算子参数可能省略默认值，同名算子也可能有不同 overload。这里复用 pnnx 已经链接的 LibTorch Dispatcher，通过 FunctionSchema 绑定参数、补默认值，并获取别名与写入信息。

这部分集中在 exported_program_operator.*，archive/JSON 解析器不使用 LibTorch 的内部 PT2 JSON serializer。pnnx 转换器仍需链接 LibTorch，生成的原生 ncnn 模型可以由 ncnn 运行时加载，部署端无需携带 LibTorch。

文件中的 ATen opset 需要与链接的 LibTorch 一致，没有实现跨 opset 的算子升级器。schema 可以读懂，不代表任意版本的算子语义都能兼容。

复用已有 lowering

PT2 导入之后使用现有 PNNX passes，没有单独复制一套 PT2 后端。部分 ATen 图形态需要补充通用转换，例如卷积、归一化、静态 weight norm 和部分循环网络算子。

共享 passes 的修改也会影响 TorchScript，因此回归同时保留两个入口。最近修复的单元素 split/chunk/unbind 就属于这类问题：PNNX Python 中保留显式解包，到原生 ncnn 后端才融合，避免把只有一个 Tensor 的 tuple 当成 Tensor 传给后续算子。

基本动态维度与再导出

已实现命名的动态输入维度、上下界、多个输入共享同一维度，以及受支持形状算子使用的运行时 aten.sym_size.int 查询。例如：

batch = torch.export.Dim("batch", min=3, max=8)
program = torch.export.export(
    model, example_inputs, dynamic_shapes=({0: batch},)
)

这些信息会保存在 PNNX IR 中，经过参数文件保存/读取后仍可用于 PT2 再导出。生成的 PNNX Python 会检查输入 rank、静态维度、动态范围和共享维度是否一致。

动态支持还有限制：2*s0、s0*s1 这样的派生表达式、数据相关形状和动态标量运算暂不支持。原生 ncnn 还受算子和 batch 布局限制，也不会执行 PNNX Python 中的 PT2 范围检查。

生成的 export_exported_program() 导出的是重建后的推理模型，返回 ExportedProgram 并保存新的 .pt2。测试会执行再导出的程序，比较多个输入尺寸下的结果。

验证与 CI

验证分成几个层次：先检查模型能否正确读入，再执行生成的 PNNX Python，与原模型比较结构、shape、dtype 和数值。再导出与原生 ncnn 执行另外验证。

算子测试复用已有模型和比较逻辑，TorchScript 与 PT2 分进程运行，并使用不同的输出文件名。测试按 pt2_frontend 和 pt2_operator 分组，可单独执行。

下面是修复提交 2c8324a5 的本地 Linux / Torch 2.13.0+cpu 验证：

检查	结果
前端、再导出和 helper	4/4 CTests 通过，包含动态维度、IR 保存/读取、空权重和错误诊断
单元素算子输出	6 种模型 × PT2/TorchScript，通过 PNNX Python 执行和 PT2 再导出比较
非标量单元素原生推理	5 种模型 × PT2/TorchScript，10 个场景的 shape 和数值通过
原生属性读取	Bool-fp32、Double-fp32、Double-fp16 三种场景通过，覆盖独立 buffer 和前后卷积权重

相邻回归首轮为 19 通过、6 失败。失败来自本机旧 ncnn Python 绑定缺少 batch_index 参数，使用兼容绑定后，这 6 项加属性测试共 7 项通过。完整 PT2 selection 包含 397 个 CTest 条目，这不代表支持 397 个独立模型，也不代表每个条目都验证了原生 ncnn 推理。最新修复后没有重新执行完整 selection，分版本记录放在 PR 正文中。

CI 工作流配置了 Ubuntu、macOS、Windows 的 Torch 2.13 构建和前端测试，Linux 算子/模型任务使用 Torch 2.12.1。本次提交的 CI尚未全部完成，跨平台结果以对应提交的实际任务为准。

编译与试用

下面以 Linux、Python 虚拟环境中的 CPU PyTorch 2.13 为例。需要 CMake 和支持 C++20 的编译器，C++20 是当前 Torch 2.13 头文件的构建要求。

git clone --branch pnnx-exported-program https://github.com/mingshi2333/ncnn.git
cd ncnn
python3 -m venv .venv
. .venv/bin/activate
python -m pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu
python -m pip install numpy

cmake -S tools/pnnx -B build-pnnx -DCMAKE_BUILD_TYPE=Release \
  -DPython3_EXECUTABLE="$(command -v python)"
cmake --build build-pnnx -j 4

把下面内容保存为 export_model.py：

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(4, 8),
    torch.nn.ReLU(),
).eval()
example_inputs = (torch.rand(3, 4),)
program = torch.export.export(model, example_inputs)
torch.export.save(program, "model.pt2")

导出、转换并再导出：

python export_model.py
./build-pnnx/src/pnnx model.pt2
python -c 'import model_pnnx; model_pnnx.export_exported_program()'

输入形状与权重从 PT2 读取，这个例子不需要额外指定 inputshape。生成的 PNNX Python 仍通过 PyTorch 执行，原生部署使用 .ncnn.param 和 .ncnn.bin。

当前范围

主要验证目标是 PyTorch 2.13 的 raw-payload PT2，代码保留部分旧 schema 的读取路径，没有持续测试所有旧版本，也没有承诺未来 schema 自动兼容。

当前不支持旧 pickle 权重布局、训练图、外部输入或 state 的修改、通用控制流，以及自定义对象等输入。原生后端也不支持所有能由 PNNX 表示的 dtype、形状和算子组合。已识别的导入限制和残留未完成 lowering 的目标会报告错误，具体范围列在 README。

后续会结合上游 review 继续整理实现，并按真实模型需要扩展动态 shape。欢迎用自己的 .pt2 模型试转换。如果遇到问题，附上 PyTorch 版本、最小导出代码和转换日志会比较方便定位。

感谢 ncnn/pnnx 和 PyTorch 提供的基础实现，也感谢 review 中对转换语义和原生输出问题的反馈。