[Tencent Rhino-Bird 2026] pnnx support for PyTorch ExportedProgram (.pt2)

[geekgoer]

pnnx support for PyTorch ExportedProgram (.pt2)

This project was developed for the 2026 Tencent Rhino-Bird Open Source Program. Its goal is to enable pnnx to directly convert models saved as PyTorch ExportedProgram packages:

pnnx model.pt2

Author: @geekgoer
Pull request: Tencent/ncnn#6959
Development branch: geekgoer/ncnn:feat/pnnx-pt2
Status: pull request under review; upstream GitHub Actions are awaiting maintainer approval
Last updated: September 12, 2026

Goals and current status

Goal	Status
Support pnnx model.pt2	Implemented
Parse the ExportedProgram graph and model state	Implemented for the supported inference scope
Avoid new third-party dependencies	Implemented with internal JSON, restricted pickle, and ZIP/ZIP64 readers
Add PT2 support to the existing pnnx tests	Implemented
Add strict invalid-package and round-trip tests	Implemented
Validate on multiple platforms	Local Windows validation completed; upstream CI is awaiting maintainer approval

Design and implementation

A PT2 package contains a serialized ExportedProgram graph, a graph signature, metadata, and tensor payloads. The new C++ frontend performs the following steps:

1. detects PT2 packages from their archive contents and validates the ZIP central directory;
2. parses the ExportedProgram JSON schema without adding a JSON dependency;
3. distinguishes user inputs, parameters, persistent and non-persistent buffers, and tensor constants using the graph signature;
4. restores tensor dtype, shape, stride, storage offset, byte order, and raw model state;
5. validates and preserves protocol-1 tuple/list input and output PyTrees;
6. constructs pnnx IR from the serialized ATen graph;
7. normalizes exported ATen forms and reuses the existing pnnx and ncnn pass pipelines.

The archive reader supports stored entries, central-directory indexing, data descriptors, ZIP64 metadata, bounds validation, and CRC checking. Compressed or encrypted entries are rejected explicitly, so no decompression library is introduced.

The generated *_pnnx.py model also provides export_exported_program(example_inputs=None). This enables the following round trip:

source PT2 -> pnnx -> generated PNNX Python -> PT2 -> pnnx

Additional normalization was added for raw ExportedProgram forms including convolution, pooling, interpolation, instance normalization, recurrent networks, channel shuffle, triangular masking, windows, and scatter/slice-copy operations.

Validation

The main local environment was Windows, MSVC Release, and PyTorch 2.5.1 CPU.

• the focused frontend, operator, PyTree, round-trip, and invalid-package tests passed;
• 389 top-level PT2 variants completed with 372 passes, 17 exact PyTorch-exporter-related skips, and 0 failures;
• optional Python dependencies used by the matrix were installed, so none of those 17 skips were missing-package skips;
• malformed JSON, duplicate keys, invalid CRC, compressed members, unsupported schemas or PyTrees, and unsupported mutations produce explicit failures;
• parameters, buffers, tensor constants, dtype/stride/storage-offset handling, and nested tuple/list inputs and outputs have focused coverage;
• schema compatibility was checked with representative schema 2, PyTorch 2.4.1 schema 5.3, PyTorch 2.5.1 schema 7.3, and PyTorch 2.13.0 packaged schema 8.20 models;
• a representative converted model matched the source ExportedProgram exactly (maxdiff = 0.0);
• 245 pnnx-to-ncnn PT2 test variants are registered; their full execution and the multi-platform matrix are pending upstream CI.

Review feedback has also resulted in focused fixes and regression coverage for:

• preserving dtype, broadcast, and copy semantics before scatter operations;
• explicitly rejecting unsupported caller-visible user-input mutations;
• preserving asymmetric height/width arguments for exported DeformConv2d;
• mapping projection weights correctly for an LSTM with bias=False and proj_size > 0;
• preserving floating-point parameter dtypes during PT2 round-trip export;
• generating valid default example tensors for boolean inputs;
• preserving requested dtype conversion when folding scalar torch.full nodes;
• using concrete example dimensions in generated TorchScript, ONNX, and PNNX export helpers for dynamic PT2 inputs;
• explicitly rejecting integer scalar constants outside the 32-bit PNNX parameter range instead of silently narrowing them;
• generating valid empty-shape example tensors for rank-0 floating-point inputs.

Supported scope

• inference-only ExportedProgram graphs with static shapes and supported named dynamic dimensions;
• tensor user inputs and outputs;
• protocol-1 tuple/list input and output containers;
• parameters, persistent and non-persistent buffers, and tensor constants represented by supported raw tensor payloads;
• ATen operators whose arguments can be represented in pnnx and whose graphs can be lowered by existing passes;
• disabled autocast and set-grad wrappers with tensor-only captured graphs.

Current limitations

• general symbolic expressions, data-dependent shapes, and symbolic scalar data flow are not supported;
• dict, namedtuple, custom PyTree containers, keyword inputs, and non-tensor input/output leaves are not supported;
• training graphs, mutation outputs, general control flow, and enabled autocast or set-grad wrappers are not supported;
• the restricted pickle reader implements only the reducers required by the validated PT2 tensor payloads;
• compressed or encrypted ZIP members are rejected;
• successful PT2-to-PNNX conversion does not imply that every operator has a native ncnn backend implementation.

Unsupported archive, schema, or graph features fail with explicit diagnostics rather than silently falling back to TorchScript.

Example

Export a model with PyTorch:

import torch

model = Model().eval()
example_inputs = (torch.rand(1, 3, 224, 224),)
program = torch.export.export(model, example_inputs)
torch.export.save(program, "model.pt2")

Convert it with pnnx:

pnnx model.pt2

Re-export the generated PNNX Python model as PT2:

python -c "import model_pnnx; model_pnnx.export_exported_program()"

Links

• Pull request: pnnx: support PyTorch ExportedProgram .pt2 models #6959
• Development branch: https://github.com/geekgoer/ncnn/tree/feat/pnnx-pt2
• PyTorch ExportedProgram documentation: https://docs.pytorch.org/docs/stable/export.html

Feedback and additional compatibility reports are welcome.