pnnx for torch exported program - #6953
Conversation
- pt2_sweep.py: 全量测试扫描, pt2/ts 双路径 .ncnn.param 结构对拍, 三分类 (PASS/DIFF/UNSUPPORTED_OP), 当前基线 PASS 117 / DIFF 8 / UNSUPPORTED_OP 90 - pt2_crosscheck.py: 算子级对拍矩阵 - testutil_pt2.py + test_pt2_smoke.py: pyncnn 数值对拍 (ctest: pt2_smoke) - test_pt2_weights.py: 权重链路验证 注: pt2_crosscheck.py 位于 tests/ncnn/ 下, 非 tools/pnnx/ 根 (CLAUDE.md 旧引用有误)
json.hpp / load_pt2.cpp(h) / load_pt2_parse.cpp(h) + main.cpp/storezip 探测分发。 作废原因 (2026-08-30 实测): torch 2.13 的 .pt2 为 JSON 图 (schema 8.20) + ZIP_STORED 原始权重, 纯 C++ 可解析, "pickle + deflate 需 Python 预处理" 的 架构前提不成立; 竞品 PR Tencent#6933 已以纯 C++ 方案占位并系统性覆盖本方案的 参数兜底/形态判别等全部设计点。 保留价值: kPt2PreprocessPy 内嵌的 aten->pnnx 语义映射与 11 条踩坑记录, 可转为对 PR Tencent#6933 的评审材料; 验证基建见上一 commit。
- 删除 load_pt2.cpp/h、load_pt2_parse.cpp/h(桥接前端,格式前提已证伪) - main.cpp / CMakeLists.txt 恢复 master(c189d88)基线 - 保留 json.hpp(自写 JSON 解析器,独立重建将复用) - 保留 storezip central directory 读取增强(通用基础设施) - 验证基建(dcedd54)与格式逆向(docs/11)原样保留 新方向:独立重建 pt2 前端,PR Tencent#6933 仅作思路参考(不抄代码)。
复用 pt2_crosscheck 的 normalize_param,直接比对已落盘的 sw_*.ncnn.param 与 sw_*_ts.ncnn.param,无需重跑 pnnx。
- pt2_schema.h/cpp:L4 schema 层,解析 model.json(header/nodes/signature) 与 weights/constants config;节点参数按 as_* 变体建模(12 种,实测 237 个 .pt2 全量普查),kind=1/2 区分 positional/keyword 实参 - load_pt2.h/cpp + main.cpp:model_file_maybe_pt2 探测(zip 签名 + <root>/models/model.json 特征),插在 torchscript 分发之前; 忠实转写 builder 为 N2 留桩 - 纯 C++ 零 libtorch,无条件编入 pnnx(pnnx_SRCS) - json.hpp 边界加固:自包含 <cstdio>、未闭合字符串报错、数字扫描收紧 (前导零/小数点/指数校验)、\u 代理对展开为 4 字节 UTF-8 - tests/test_json.cpp:70 项边界单测(独立 harness); tests/test_pt2_schema.cpp:schema dump harness(-c 规范化输出) - 验收:239 个真实 .pt2 双侧规范化对拍全一致(C++ harness vs pt2-dump/dump_canonical.py);WSL 构建通过,pnnx mini.pt2 正确命中 pt2 分发
- load_pt2 builder:忠实转写,与 ts level1 通用转写形态对齐—— 权重/buffer/tensor_constant → pnnx.Attribute(op 名 = state_dict 名, 裸字节从 zip 读入);标量参数 → prim::Constant operand;张量列表 → prim::ListConstruct;Input/Output 命名与 ts 相同;op 保持 aten 原名 (target 去 torch.ops. 前缀与 overload 后缀);零归一化 - pass_level2/F_pt2.cpp:PT2 形态分支,把 torch.export 的缺省实参形态 归一成 ts 等价形态——flatten 2-op 补 end_dim=-1 汇合 torch_flatten; conv2d 四种缺省/全参变体 emit params 化 F.conv2d(缺省 dilation/groups 内联,N3 换静态表)。priority 55:先于 torch_flatten(60)归一形态, 先于 F_conv2d_1(140)避免其 7-input pattern 误吃 - schema:补 constants 路径(tensor_N 位于 data/constants/)与 zippath - 验收(WSL,同权重模型):conv2d 模型 ts/pt2 双路径 .ncnn.param 与 .ncnn.bin 全部逐字节一致(权重字节链路端到端打通);smoke 模型 (flatten+cat+relu)param/bin 逐字节一致;aten::add 2-op 与 cat 列表 形态由现有 fuse_expression/torch_cat 直接消化,零新 pass
- scripts/dump_aten_defaults.py:离线 dump torch._C._jit_get_all_schemas, --scan 语料驱动按需收录,180 算子入库(可审计、可再生成、零运行时依赖) - src/aten_defaults_table.h:生成的静态表(header-only,按 overload 全名查表) - load_pt2:节点缺参查表补全为完整 schema 形态(带 fill default 标记), provided 形参名与表不符时回退原样转写;F_pt2.cpp 的 5 条内联缺省分支删除, pt2 图与 ts 图同构后由既有 ts 形态分支零改动消费 - pt2_schema:解析 graph.tensor_values 张量元数据表(含中间张量形状/dtype, 修正"JSON 不携带中间张量形状"的误记);builder 填 operand.shape, 文件事实优先于 CLI inputshape - 空列表默认值编为 type 0(None)对齐 ts 形态(max_pool2d stride=() 段错误根因, 修复后 sweep PNNX_PT2_FAIL 120 -> 6)
- pt2_schema: 新增 json_as_int 助手替换全部裸 asInt(),整数字段被未来 torch 版本写成浮点时不再经 JSON_INT/DOUBLE union 静默错值; schema_version major/minor 带 isNumber 守卫,保持缺失=-1 语义 - load_pt2: 删除与 3.5 节 tensor_values 回填矛盾的过期注释(N3 漏改) - dump_aten_defaults: 空列表编码注释改为实际机制说明(编 INTS 空串, builder 转 None 对齐 ts 形态) - 验证:WSL 重建后 test_pt2_schema mini/smoke_cat/mega 解析 OK, 全量 sweep 复跑 PASS 115 / DIFF 95 / PNNX_PT2_FAIL 6 / 219 与基线 逐项一致(零回归)
…S 115→178 builder(load_pt2): - hoist_constants:标量 prim::Constant 前移到消费者之前。builder 惰性创建 的常量滞后于消费者,fuse_expression(level3)反向扫描会先把常量包成 pnnx.Expression,算术链融合无法内联字面量,产出与 ts 不同的常量 blob 形态(激活族等 36 个 DIFF 的根因) - 切分族(unbind/split/split_with_sizes/chunk/tensor_split)转写为 1 输出 + prim::ListUnpack,对齐 ts level1 形态,复用现有 torch_* 形态 分支与 fuse_op1ton_unpack(level3)折叠(getitem 已被 exporter 折叠成 as_tensors 多输出,1 输出 pattern 无法匹配) pass_level2 PT2 形态分支: - F_conv1d_1/F_conv3d_1/F_conv_transpose1d_1/2d_1/3d_1:torch.export 产出 公开算子 aten::conv1d/conv3d/conv_transpose*(与 ts 内部算子 aten::_convolution 不同),参数实参经 fuse_constant_expression 折入 params,后续 fuse_static_conv* 正常折叠权重(F_conv2d_1 同款先例) - F_pt2_weight_norm:aten::_weight_norm(v,g,dim=0) 折成 pnnx.Attribute, 复用 utils.cpp apply_weight_norm(与 ts level1 同一 float 实现,权重 字节一致);v/g 非常量等形态不匹配显式失败 - F_pt2_adaptive_pool_*:adaptive pool output_size 中被 torch.export 实例 化的 None(=输入空间维)还原为 0(pass_ncnn 对 0 写 -233 哨兵,恒等池化 语义不变);仅匹配 prim::Constant 形态,替换图同构靠'确会改写才匹配' 终止重写循环 验证:全量 sweep PASS 178 / DIFF 32 / PNNX_PT2_FAIL 6 / 219(基线 115/95/6),0 回归 63 改善;F_relu/stack 等代表性场景 param+bin 逐字节 一致;clang-format 10.0.1 已跑(astyle 同前,提 PR 前补验)
builder 解析 node.metadata.nn_module_stack(export 图一等事实),对白名单 (模块类, aten 算子)对转写为 nn.<类> + params 化形态——折参规则逐模块对齐 ts 侧 level1 模块转换(FuseModulePass)的产出,与 F. 形态对齐 level1 通用 转写同理,loader 侧转写与 ts 层次对等。 - 白名单:ReLU6/Softmax2d/ChannelShuffle/PixelShuffle/MaxPool1d-3d(含 with_indices)/AdaptiveAvgPool1d-3d/pad 族 9 类/Upsample 三类 - 折参改名对齐 level1:pad→padding、output_size→size、scale_factors→ scale_factor;折参集合精确对齐(pnnx 模式匹配对参数集合敏感) - 单值 int/float 实参按算子空间维广播成列表(对齐 JIT 的 schema 泛化) - export 省略的默认实参按 L7 表补进 params(MaxPool 的 dilation/ceil_mode) - F_pt2 新增 nn.AdaptiveAvgPool* params 形态的 None 实例化还原分支 - MaxPool with_indices → nn.MaxPool* + return_indices=True,indices 由 eliminate_maxpool_indices 消除后与 ts 形态汇合 sweep:PASS 197 / DIFF 13 / PNNX_PT2_FAIL 6 / 219(基线 178/32/6)
…7→204,PNNX_PT2_FAIL 清零 - F_pt2_fold_ones_like:torch.ones_like+add(标量) 静态折成常量 Attribute。 ts 侧靠 pass_level0 跑 libtorch 折常量子图;pt2 零 libtorch,利用 ones_like 值语义恒为全 1 静态求值,匹配 torch_ones_like(priority 20)归一后的形态, 修复 maximum/minimum/atan2/pow 的 ones_like(z)+0.5 常量族(4 场景) - LayerNorm/RMSNorm 模块形态:折 normalized_shape/eps,elementwise_affine 由 weight 实参存在性判定;γ/β 折 op attrs(ts level1 模块转换同构, pass_ncnn 按 @weight/@bias 捕获)(2 场景) - argument_to_constant 补 MEMORY_FORMAT/DEVICE 枚举实参转写(clone/stft/ istft 的 memory_format/device 实参,PNNX_PT2_FAIL -1) - testutil_pt2 数值对拍语义修复:batch_index=233 的 3D size-1 输入先剥 batch 维喂 ncnn;batch_index=0 输出还原形状再比。权重数值对拍 7/7 全绿 (fp16 存储下 max|d|≈3e-4),权重字节链路 state_dict→zip→Attribute→ .bin→推理端到端验证 sweep:PASS 204 / DIFF 12 / PNNX_PT2_FAIL 0 / EXPORT_FAIL 1 / SKIP 2 / 219
…0 不变 - F_pt2_fold_ones_like:输出 dtype/静态 shape/溢出校验前移到 match(), 非 f32、缺 shape、非正维、非标量 other 保持原图,write() 不再静默 留下无 data 的空 pnnx.Attribute - load_pt2 argument_to_constant:DEVICE 保留 index 编码为 type:index, cuda:1 不再降级为 cuda;无 index 为裸 type;空 device 编 None - testutil_pt2._restore_ncnn_output:仅允许剥 size-1 batch 轴这一种 还原关系(batch_index=0/233 统一语义),错序但 numel 相同的形状拒绝; 修复上一轮 233 分支过度收紧误杀 weights 场景合法剥离 - 新增 test_pt2_regress.cpp:ones_like 5 形态 + DEVICE 5 编码白盒单测 (零 libtorch,include 产品 cpp,链接行需置于最后避开注册表静态 初始化顺序问题) - 新增 test_pt2_testutil.py:helper 剥离/拒绝语义 9 断言 验证:全量 sweep PASS 204 / DIFF 12 / PNNX_PT2_FAIL 0 / TOTAL 219 与 N4 基线一致;ctest pt2 smoke 1/1;数值对拍 7/7;clang-format 10.0.1 dry-run 零 violation
新增 DIFF 基线夹具,sweep PASS 204→208 - 新增 tests/ncnn/pt2_diff_fixture.py:12 个 DIFF 场景基线夹具 (capture/verify 双模式),落盘 pnnx IR/ncnn param/bin/源模型供 M3-M5 改造后快速回归 - fuse_static_conv 新增 6 个 *_pad 融合 pass(1/2/3D × bias 有无, 先于静态折叠):F.pad(reflect/replicate) + F.conv*d(zeros) → nn.Conv*d(padding_mode=...,padding 折算),与 ts level1 nn_Conv* 同构,收敛 Conv3d reflect/replicate + groups 场景 - weight_norm 改判命令式 fold_pt2_weight_norm(pass_level2 主函数 fuse_constantlist 后调用):v/g 均为 pnnx.Attribute 且 dim=0 时就地 转 Attribute,绕开 pattern 引擎"消费者数==anchor"约束,收敛 weight_norm→Linear 多消费者场景 - F_pt2 新增 F_pt2_local_response_norm_base + nn.LocalResponseNorm / F.local_response_norm 两子类:pt2 静态 shape(含 -1)对齐 ts 动态 size+ListConstruct,收敛 2 个 LRN 场景 - Upsample/interpolate 含 .5 scale 场景定性「格式级不可达」:export 将 floor(size×scale) 子图内联为常量且 floor 有损,scale 不可恢复 验证:全量 sweep PASS 208 / DIFF 8 / EXPORT_FAIL 1 / SKIP 2 / PNNX_PT2_FAIL 0 / TOTAL 219;pyncnn 数值对拍 7/7;ctest pt2 smoke 1/1;clang-format 10.0.1 dry-run 对 4 个改动文件零 violation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 98f5499f1c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (fread((char*)&z64_cd_records, sizeof(z64_cd_records), 1, fp) != 1) | ||
| continue; | ||
| if (fread((char*)&z64_cd_size, sizeof(z64_cd_size), 1, fp) != 1) | ||
| continue; | ||
| if (fread((char*)&z64_cd_offset, sizeof(z64_cd_offset), 1, fp) != 1) |
There was a problem hiding this comment.
Consume the total-record count in Zip64 EOCD
After the fixed fields, a Zip64 EOCD contains four consecutive 64-bit values: records-on-disk, total-records, central-directory size, and central-directory offset. These reads consume only three, so z64_cd_size receives the total-record count and z64_cd_offset receives the directory size; cd_offset_valid() consequently fails. This breaks every archive emitted by this file's StoreZipWriter, which always writes a Zip64 EOCD, preventing pnnx from reopening its own parameter and foldable-constant archives.
Useful? React with 👍 / 👎.
| if (dim_index >= 0 && dim_index < (int)ishape.size() && osz.ai[i] == ishape[dim_index]) | ||
| { | ||
| osz.ai[i] = 0; |
There was a problem hiding this comment.
Preserve explicit adaptive-pool output sizes
When an explicit adaptive-pool size happens to equal the traced input dimension, this treats the equality as proof that the source value was None and rewrites the literal to 0. The ncnn conversion maps 0 to the preserve-input-dimension sentinel, so a model traced with an 8×8 input and explicitly requesting (8, 8) incorrectly returns 10×10 when later run on 10×10 input instead of pooling to 8×8. Because this pass is globally registered, the regression also affects existing TorchScript graphs, not only PT2 inputs.
Useful? React with 👍 / 👎.
| if (spec.kind == Pt2InputSpec::TENSOR_CONSTANT) | ||
| entry = program.find_constant(spec.state_dict_name); | ||
| else | ||
| entry = program.find_weight(spec.state_dict_name); |
There was a problem hiding this comment.
Load non-persistent buffers from constants
For an exported module using register_buffer(..., persistent=False), the graph signature still identifies the input as BUFFER, but its bytes reside in the exported program's constants rather than its state-dict weights. This branch sends every buffer to find_weight, and the later call also selects the weights path, so conversion aborts with "weight entry not found" for such models. Treat non-persistent buffers as constants for both lookup and archive-path selection.
Useful? React with 👍 / 👎.
| Pt2OutputSpec s; | ||
| s.graph_name = parse_spec_graph_name(it->second); | ||
| out.push_back(s); |
There was a problem hiding this comment.
Filter mutation specs from public outputs
When the exported signature contains buffer_mutation or user_input_mutation output specs, this loop appends them exactly like user_output without checking the spec kind. load_pt2 subsequently creates a public pnnx.Output for each mutation value, changing the model's output count and ordering while not implementing the mutation semantics. Restrict this list to user_output entries, or reject unsupported mutation-bearing exports explicitly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f15074dffd
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| pnnx_ncnn_add_test(pt2_smoke) | ||
| pnnx_ncnn_add_test(pt2_weights) |
There was a problem hiding this comment.
Exclude PT2 tests from legacy Torch jobs
These tests are registered unconditionally, so the existing workflow's bare ctest --output-on-failure -j 8 also runs them in every legacy matrix entry, including Torch 1.8–1.13 where torch.export does not exist. CTest documents -R as “Run tests matching regular expression”; because that job has no selector or exclusion, run_pt2_test() catches the resulting export error and returns false, causing at least test_ncnn_pt2_smoke and test_ncnn_pt2_weights to fail every affected matrix job. Gate these registrations/tests by Torch capability or exclude them from the legacy run while retaining the dedicated pt2-test job.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 39dc6a04a0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (aten_type == "aten::adaptive_avg_pool1d" || aten_type == "aten::adaptive_avg_pool2d" | ||
| || aten_type == "aten::adaptive_avg_pool3d" || aten_type == "aten::adaptive_max_pool1d" | ||
| || aten_type == "aten::adaptive_max_pool2d" || aten_type == "aten::adaptive_max_pool3d") | ||
| { | ||
| op->name = "pt2_" + op->name; |
There was a problem hiding this comment.
Preserve explicit adaptive-pool output sizes
The new source guard does not resolve the prior issue because the loader assigns the pt2_ marker to every exported adaptive-pool node, regardless of whether output_size originally contained None. Consequently, an explicitly requested size equal to the export input dimensions still passes the guard in F_pt2_adaptive_pool_base and is rewritten to the preserve-dimension sentinel; for example, exporting explicit (8, 8) at 8×8 and later running at 10×10 produces 10×10 instead of 8×8.
Useful? React with 👍 / 👎.
| fseek(fp, extra_size - 4, SEEK_CUR); | ||
| extra_offset += extra_size; |
There was a problem hiding this comment.
Skip the full payload of preceding ZIP extra fields
For a valid stored Zip64 archive where another extra field precedes the 0x0001 Zip64 field (for example an extended-timestamp field), extra_size is the payload length after the four-byte ID/length header. Seeking by extra_size - 4 and advancing extra_offset by only extra_size leaves the cursor inside that payload, so subsequent bytes are interpreted as another header; depending on those bytes, opening the archive can fail or loop indefinitely before pnnx even detects the PT2 model. Skip all extra_size payload bytes and account for the four-byte header in the offset.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ebd45082c5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ) | ||
| output_count = len(re.findall(r"^pnnx.Output\s", param, re.MULTILINE)) | ||
| check( | ||
| result.returncode == 0 and output_count == 1, |
There was a problem hiding this comment.
Align the mutation regression with explicit rejection
For the BufferMutation model, the exported signature contains a buffer_mutation output, and the current parse_output_specs() now deliberately returns -1 for every non-user_output spec. Therefore this assertion can never pass: _export_and_convert() receives a nonzero return code and no valid parameter output. The dedicated pt2-tests workflow runs test_ncnn_pt2_review, so update the regression to expect rejection or change the parser to filter mutations while succeeding.
Useful? React with 👍 / 👎.
| if (extra_size < sizeof(zip64_extended_extra_field)) | ||
| { | ||
| // skip this extra field block | ||
| fseek(fp, extra_size - 4, SEEK_CUR); | ||
| fseek(fp, extra_size, SEEK_CUR); | ||
| extra_offset += extra_size; | ||
| continue; |
There was a problem hiding this comment.
Parse only the Zip64 fields that are present
When a standard Zip64 central-directory entry saturates only its sizes or its sizes plus local-header offset, the 0x0001 payload is respectively 16 or 24 bytes; fields are included conditionally in the order of the saturated 32-bit fields. Requiring the full 28-byte struct skips these valid payloads, leaving the sizes at 0xffffffff, so archives containing an entry over 4 GiB are recorded with the wrong size and subsequently fail or attempt a multi-gigabyte allocation. Read each required field individually based on which legacy field is saturated.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8d8f4ac646
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (dim_index >= 0 && dim_index < (int)ishape.size() && osz.ai[i] == ishape[dim_index]) | ||
| { | ||
| osz.ai[i] = 0; |
There was a problem hiding this comment.
Preserve each None axis in adaptive pooling
For an export at 8×8 with output_size=(None, 8), both serialized values are 8, and this loop rewrites both to the preserve-input sentinel, so inference at 10×10 produces 10×10 instead of 10×8. Fresh evidence after the earlier guard change is that adaptive_pool_has_none records only whether any axis was None, losing which axis it was; preserve the per-axis mask rather than rewriting every value equal to the traced input size.
Useful? React with 👍 / 👎.
| return -1; | ||
| } | ||
|
|
||
| attr.data = raw; |
There was a problem hiding this comment.
Materialize strided PT2 weights before storing attributes
When a parameter or buffer is non-contiguous, such as a transposed 2-D weight, the raw entry contains storage-order bytes while entry.sizes, entry.strides, and entry.storage_offset describe its logical view. Assigning those bytes directly to a contiguous pnnx attribute silently permutes values when the storage size equals the tensor's element count, and offset views instead fail the preceding size check; materialize the logical contiguous tensor using the serialized strides and offset.
Useful? React with 👍 / 👎.
| if (z64_total_cd_records != z64_cd_records) | ||
| continue; | ||
|
|
||
| if (cd_offset_valid(z64_cd_offset) && z64_cd_offset + z64_cd_size <= (uint64_t)file_size) |
There was a problem hiding this comment.
Accept zero-entry Zip64 archives
When Graph::save serializes a weightless graph, StoreZipWriter emits a valid Zip64 archive with zero records, a zero-size central directory, and z64_cd_offset == 0. This condition nevertheless requires a central-directory file-header signature at offset zero, so StoreZipReader::open rejects the generated .pnnx.bin; handle the zero-record/zero-size directory as a valid special case.
Useful? React with 👍 / 👎.
| if (eocd_cd_offset != 0xffffffff) | ||
| { | ||
| zip64_end_of_central_directory_locator eocdl64; | ||
| fread((char*)&eocdl64, sizeof(eocdl64), 1, fp); | ||
| if (cd_offset_valid(eocd_cd_offset) && (uint64_t)eocd_cd_offset + eocd_cd_size <= (uint64_t)file_size) | ||
| { | ||
| cd_offset = eocd_cd_offset; | ||
| cd_size = eocd_cd_size; | ||
| cd_records = eocd_records; |
There was a problem hiding this comment.
Read Zip64 counts when the EOCD count saturates
For a valid Zip64 archive with 65,536 or more entries but a central-directory offset below 4 GiB, the legacy EOCD offset remains finite while its record-count fields are 0xffff. This branch therefore skips the Zip64 locator and scans only 65,535 entries, silently omitting the rest from filemetas; select the Zip64 record when any legacy EOCD field is saturated, not only the offset.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 90942a801b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (dim_index >= 0 && dim_index < (int)ishape.size() && osz.ai[i] == ishape[dim_index]) | ||
| { | ||
| osz.ai[i] = 0; |
There was a problem hiding this comment.
Preserve explicit sizes in adaptive-pool modules
For nn.AdaptiveAvgPool* with an explicit size equal to the traced input dimension, this separate module-form pass still infers None from equality and rewrites that size to 0. Fresh evidence after the function-form marker fix is that this matcher has no __pt2_none_axes or PT2-source guard, so it also rewrites existing TorchScript module graphs; for example, nn.AdaptiveAvgPool2d((8, 8)) traced at 8×8 becomes preserve-both-dimensions and produces 10×10 rather than 8×8 on a 10×10 input.
Useful? React with 👍 / 👎.
| || node.target.find("adaptive_max_pool") != std::string::npos) | ||
| && node.stack_trace.find("output_size=") != std::string::npos) |
There was a problem hiding this comment.
Detect positional None axes in adaptive pooling
For a normal positional call such as F.adaptive_avg_pool2d(x, (None, 8)), the recorded source line does not contain the literal output_size=, so this condition never builds the per-axis marker. Fresh evidence after the earlier mask fix is that marker discovery is still restricted to keyword spelling; the serialized concrete (8, 8) is consequently retained, and running the converted model on a 10×10 input yields 8×8 instead of the required 10×8.
Useful? React with 👍 / 👎.
| fprintf(stderr, "load_pt2_schema: unknown argument variant:"); | ||
| for (std::map<std::string, JsonValue>::const_iterator it = arg.object_value.begin(); | ||
| it != arg.object_value.end(); ++it) | ||
| { | ||
| fprintf(stderr, " %s", it->first.c_str()); | ||
| } | ||
| fprintf(stderr, "\n"); | ||
| return Pt2Argument::NONE; |
There was a problem hiding this comment.
Reject symbolic arguments instead of converting them to None
When a dynamic exported graph contains the standard as_sym_int or as_sym_ints argument variants, such as a reshape size derived from an input dimension, they fall through this detector and are returned as NONE. load_pt2 then emits a valid prim::Constant(None) rather than rejecting or preserving the symbolic value, so conversion can succeed with incorrect operator arguments; unsupported symbolic variants should at least fail explicitly instead of being conflated with as_none.
Useful? React with 👍 / 👎.
| if (node.inputs[j].name == "weight" && node.inputs[j].arg.type == Pt2Argument::TENSOR) | ||
| { | ||
| has_weight = true; | ||
| break; | ||
| } | ||
| } | ||
| op->params["elementwise_affine"] = has_weight; |
There was a problem hiding this comment.
Handle LayerNorm modules with weight but no bias
For nn.LayerNorm(..., elementwise_affine=True, bias=False), export supplies a tensor weight and a None bias, but this marks the module affine solely from the weight. The resulting nn.LayerNorm has elementwise_affine=True and @weight but no @bias, while the ncnn lowering pattern in pass_ncnn/nn_LayerNorm.cpp requires both attributes, so conversion cannot lower this valid module; synthesize the zero bias or preserve a weight-only representation.
Useful? React with 👍 / 👎.
| if (raw.size() != storage_count * elemsize) | ||
| { | ||
| fprintf(stderr, "load_pt2: weight size mismatch %s: expect %zu got %llu\n", entry_path.c_str(), | ||
| storage_count * elemsize, (unsigned long long)raw.size()); | ||
| return -1; |
There was a problem hiding this comment.
Accept backing storage larger than the tensor span
For a view-backed parameter whose underlying storage contains unused trailing elements, the raw PT2 entry contains the complete backing storage while storage_count is only the highest referenced element plus one. The new stride materialization therefore still rejects valid views at this exact-equality check—for example, a contiguous base[:3] parameter backed by a 10-element storage expects 3 elements here but reads 10—even though every logical index is in bounds; validate that the raw storage is at least the required span and materialize from it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 49b92778e4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if ((node.target.find("adaptive_avg_pool") != std::string::npos | ||
| || node.target.find("adaptive_max_pool") != std::string::npos) | ||
| && node.stack_trace.find("None") != std::string::npos) |
There was a problem hiding this comment.
Preserve None axes when the size is not inline
When an adaptive-pool size containing None is stored in a variable or module attribute—for example the existing nn.AdaptiveAvgPool2d((None, 3)) pattern—the exported stack trace contains an expression such as self.pool_2(x) or self.output_size, not the literal text None. This condition therefore never records __pt2_none_axes; the concrete export-time height remains fixed, so running the converted model with a different height produces a fixed-height output instead of preserving that input dimension. The axis mask needs to come from data that survives export rather than requiring None to appear in source text.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b7512b2d4b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| case Pt2Argument::SCALAR_TYPE: | ||
| case Pt2Argument::MEMORY_FORMAT: | ||
| // dtype / memory_format 枚举按整数值转写(与 ts 侧常量物化一致) | ||
| value = Parameter((long long)a.int_value); |
There was a problem hiding this comment.
Translate serialized scalar-type enum values
PT2's as_scalar_type uses the same serialized enum demonstrated above for tensor metadata (7 is float32 and 5 is int64), but the existing level-2 passes interpret JIT scalar-type values (6 as float32, 4 as int64). Copying the value unchanged therefore shifts every explicit dtype: for example, exporting torch.ones_like(x, dtype=torch.float32) records 7, which torch_ones_like converts to torch.double, producing the wrong output dtype or an unsupported lowering. Convert the serialized scalar-type enum to the JIT/pnnx representation before creating the constant.
Useful? React with 👍 / 👎.
| adaptive_pool_none_axes.push_back(dim_index >= 0 && dim_index < (int)input_shape.size() | ||
| && output_size[j] == input_shape[dim_index]); |
There was a problem hiding this comment.
Do not infer module None axes from equal dimensions
For nn.AdaptiveAvgPool2d((8, 3)) exported with an 8-pixel input height, the module's internal stack trace also contains self.output_size, so this equality test marks the explicit first axis as None; the later pass rewrites it to the preserve-input sentinel, and inference on a 10-pixel height returns 10×3 instead of 8×3. Fresh evidence beyond the earlier finding is the newly added self.output_size fallback, which is shared by explicit and None module configurations and therefore cannot distinguish them.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 631bc3ada3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| default: | ||
| return scalar_type; |
There was a problem hiding this comment.
Translate every serialized scalar-type value
When an exported operator specifies a dtype other than the two special-cased types, this fallback silently preserves the PT2 enum even though the existing level-2 passes expect JIT enum values; for example, serialized float16 is 6, but torch_ones_like.cpp interprets 6 as torch.float, changing a requested float16 result to float32. Fresh evidence after the earlier dtype fix is that only serialized values 7 and 5 are translated; the complete PT2 scalar-type mapping is still required.
Useful? React with 👍 / 👎.
| if (eocd_records == 0 && eocd_cd_size == 0 && eocd_cd_offset == 0) | ||
| { | ||
| cd_offset = 0; | ||
| cd_size = 0; | ||
| cd_records = 0; |
There was a problem hiding this comment.
Validate the EOCD comment length before accepting a candidate
When a valid nonempty ZIP comment contains PK\x05\x06 followed by zero-valued EOCD fields, the backward scan reaches that comment sequence first and this branch accepts it as an empty archive without checking that the candidate plus its declared comment length reaches EOF. The reader consequently exposes no entries (and PT2 detection falls through to the wrong loader); verify each candidate's comment length against the file end before accepting the zero-entry special case.
Useful? React with 👍 / 👎.
| long lp = loc_pos - scan_start; | ||
| if (!(buf[lp] == 0x50 && buf[lp + 1] == 0x4b && buf[lp + 2] == 0x06 && buf[lp + 3] == 0x07)) |
There was a problem hiding this comment.
Read Zip64 locators that precede the EOCD scan buffer
For a valid Zip64 archive with an EOCD comment of 65,516–65,535 bytes, the EOCD candidate lies within the first 20 bytes of this buffer, so its locator is before scan_start and lp becomes negative. Indexing buf[lp] is undefined behavior and in practice rejects such archives; either include the extra 20 locator bytes in the scan buffer or read the locator directly from loc_pos.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 824a049658
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| case 7: | ||
| return 1; // f32 | ||
| case 5: | ||
| return 5; // i64 | ||
| default: | ||
| return 0; |
There was a problem hiding this comment.
Map all supported PT2 tensor dtypes
For an exported model with a float16 parameter or buffer, PT2 records dtype enum 6, but this mapping returns 0; load_weight_attribute() then aborts with unsupported weight dtype 6. This prevents conversion of common half-precision models even though pnnx already supports f16 attributes, and the same omission affects other supported types such as f64, i32, i16, i8, u8, and bool.
Useful? React with 👍 / 👎.
| case Pt2Argument::MEMORY_FORMAT: | ||
| value = Parameter((long long)a.int_value); | ||
| return true; |
There was a problem hiding this comment.
Translate serialized memory-format enums
When an exported operator explicitly requests torch.contiguous_format, PT2 serializes the memory-format enum as 1 because 0 is the serialization sentinel, while the existing TorchScript passes interpret 0 as contiguous and 1 as preserve. Passing the value through unchanged therefore turns an explicitly contiguous clone of a non-contiguous tensor into a preserve-format clone, which can change strides and break a subsequent operation such as view; translate this enum just as the scalar-type branch does.
Useful? React with 👍 / 👎.
the torch.export archive layout drifts across torch releases; 2.9 is the lowest version emitting a parseable archive. disable instead of fail so the version matrix (1.8-2.8) stays green.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c2a20b13a3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| names.push_back(v[i]["name"].asString()); | ||
| } |
There was a problem hiding this comment.
Preserve null slots in tensor argument lists
When an exported operator has a Tensor?[] argument, null entries are positional and cannot be discarded. For example, aten::index.Tensor represents x[:, index] with an indices list containing a null slot before the index tensor; this loop keeps only object-valued entries, so the generated prim::ListConstruct becomes [index] and applies the index to dimension 0 instead of dimension 1. Preserve each null entry as a None operand rather than compacting the list.
Useful? React with 👍 / 👎.
| if (r->shape.empty()) | ||
| { | ||
| for (size_t j = 0; j < input_shapes[input_index].size(); j++) | ||
| r->shape.push_back((int)input_shapes[input_index][j]); | ||
| } |
There was a problem hiding this comment.
Apply concrete input shapes to symbolic PT2 dimensions
For a dynamically exported PT2 model, parse_tensor_values() records symbolic dimensions as -1, making r->shape nonempty. Even when the caller supplies a concrete inputshape, this guard therefore refuses to apply it, leaving symbolic dimensions in the graph and causing shape-dependent normalization and batch-layout passes to operate on -1 values. Fill symbolic dimensions from the supplied shape (while retaining fixed exported dimensions) rather than using the override only when the entire shape is absent.
Useful? React with 👍 / 👎.
Motivation
torch.exportis PyTorch's modern export path and produces.pt2exported programs.pnnx currently supports TorchScript archives (
.pt) but cannot consume.pt2filesdirectly.
This PR adds native
.pt2loading and conversion support.Design
The implementation follows a clear separation:
.pt2loader faithfully transcribes the exported graph.pass_level2.pnnx.Attributeoperands.graph.tensor_valuesprovides authoritative tensor shape and dtype metadata.The implementation includes:
.pt2file detection and dispatchtorch.exportand TorchScript dual-path comparison helpersVerification
Structural parity sweep:
Numeric parity:
Regression tests:
test_ncnn_pt2_smoketest_ncnn_pt2_weightstest_ncnn_pt2_testutilA dedicated PyTorch 2.13 PT2 CI job has been added. The existing TorchScript test
matrix remains unchanged.
Known limitations
The remaining 12 structural differences are explicitly classified as:
Additional current boundaries:
dim=0ones_likefolding is restricted to validated static f32 forms