pnnx: support PyTorch ExportedProgram .pt2 models

[geekgoer]

Summary

• add direct pnnx model.pt2 support for PyTorch ExportedProgram archives
• parse the serialized graph, graph signature, tensor metadata, state dict, buffers, and tensor constants without adding third-party dependencies
• support legacy and packaged PT2 layouts, ZIP64 archives, and exported ATen graph normalization
• support dynamic tensor dimensions used by exported shape operations
• preserve protocol-1 tuple/list input and output PyTrees
• generate an export_exported_program() helper for PT2 round-trip export
• reject malformed, compressed, corrupt, unsupported, or state-mutating packages with explicit diagnostics
• add focused ExportedProgram tests and PT2 variants for the existing pnnx and pnnx-to-ncnn test suites
• document how to export and convert a .pt2 model

Usage

Export a model with PyTorch:

import torch

model = Model().eval()
inputs = (torch.rand(1, 3, 224, 224),)

exported_program = torch.export.export(model, inputs)
torch.export.save(exported_program, "model.pt2")

Convert the exported program with pnnx:

pnnx model.pt2

The generated PNNX Python model can also be exported to another PT2 package:

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

An optional input tuple can be passed to export_exported_program(example_inputs=None) to override the generated example inputs.

Implementation

The importer contains a schema-oriented JSON parser and a restricted pickle reader for tensor payloads. It reconstructs PNNX IR directly from the serialized ExportedProgram graph and graph signature.

The existing store-only ZIP implementation is extended to read central directories, data descriptors, ZIP64 offsets, and large tensor files. Archive entry metadata, bounds, duplicate names, encryption state, and CRC are validated before use. Compressed entries are detected and rejected explicitly without adding a decompression dependency.

Parameters, persistent buffers, non-persistent buffers, and tensor constants are restored as PNNX attributes instead of being exposed as runtime inputs. Tensor dtype, shape, stride, storage offset, byte order, and raw storage data are handled by the importer.

Additional graph-rewrite patterns normalize decomposed ExportedProgram graphs for convolution, pooling, interpolation, normalization, RNN, window, scatter/slice-copy, ChannelShuffle, torch.tril, and selected torchvision operators.

The generated *_pnnx.py model validates and reconstructs supported tuple/list input and output structures and provides export_exported_program() for a PT2-to-PNNX-to-PT2 round trip.

No new third-party dependency is introduced.

Tests

Local environment: Windows, MSVC Release, PyTorch 2.5.1 CPU.

Focused tests:

• test_exported_program: passed
• test_exported_program_ops: passed, including a dynamic-shape reshape case
• test_exported_program_pytree_roundtrip: passed, including nested inputs and outputs and PT2 -> PNNX Python -> PT2 -> pnnx conversion
• test_exported_program_invalid: passed, including archive, JSON, schema, PyTree, CRC, compression, and mutation diagnostics
• existing TorchScript regression tests for ReLU, functional Conv2d, and nn.Conv2d: passed

PT2 test matrix:

• 389 top-level PT2 variants
• 372 passed
• 17 skipped because of PyTorch 2.5 exporter limitations
• 0 failed
• all optional Python dependencies used by this matrix were installed; none of the 17 skips were caused by missing packages

Compatibility samples were checked with:

• legacy schema major 2
• PyTorch 2.4.1, schema 5.3
• PyTorch 2.5.1, schema 7.3
• PyTorch 2.13.0 packaged PT2, schema 8.20

A representative ExportedProgram and the generated PNNX Python model produced an exact match:

maxdiff = 0.0

A total of 245 pnnx-to-ncnn PT2 test variants are registered. Their full execution and the multi-platform/version matrix are pending upstream CI.

Current limitations

• user input and output leaves must be tensors
• protocol-1 tuple/list input and output PyTrees are supported; dict, namedtuple, custom containers, and keyword user inputs are not supported
• the pickle reader intentionally implements only the reducers required by supported PT2 tensor payloads
• compressed or encrypted ZIP entries are rejected
• enabled autocast/set-grad wrappers, general control flow, and unsupported mutation outputs are rejected
• user-input mutations are rejected because caller-visible mutation side effects are not currently reproduced
• successful PT2-to-PNNX conversion does not guarantee that every converted operator has a native ncnn backend implementation

Related issue

Related to #6366.

This work was developed as part of the 2026 Tencent Rhino-Bird Open Source Program.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1ac380c8cd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Reject or implement user-input mutations

For an ExportedProgram whose signature contains user_input_mutation, this silently drops the mutation output and still reports successful conversion. Calling the generated model therefore does not update the caller's tensor as the exported module does; either emit the corresponding copy-back operation or reject this output kind, as is already done for unsupported buffer mutations.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cd0dbbcc87

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".