PVN3D ONNX 转换与测试记录

1. 目的

本文单独记录当前 pvn3d-dev 容器内的 ONNX 转换流程、测试方法，以及已经确认的导出问题。

这篇文档只关注两件事：

1. 如何把当前 deploy/ 里的可导出子图转成 ONNX
2. 导出过程中实际遇到了什么问题，如何判断，下一步怎么处理

不在本文展开的内容：

• TensorRT 安装
• 容器创建
• PointNet2 插件化

这里先说明一个边界：

• 当前可以用 ONNX Runtime 测试整条 PVN3D 推理链
• 但不是“整网纯 ONNX Runtime”
• 当前可行的是“ONNX Runtime + PointNet2 Native CUDA”的混合推理流程

原因很简单：

• rgb_backbone 可以走 ONNX Runtime
• fusion_head 可以走 ONNX Runtime
• PointNet2 仍然保留 PyTorch CUDA 扩展

相关文档：

• pvn3d_tensorrt_env_setup_guide.md
• tensorrt_install_in_pvn3d_dev.md

2. 当前容器环境

本文基于当前实测容器环境编写：

• 容器：pvn3d-dev
• 仓库路径：/workspace/workflow/self/PVN3D
• Conda 环境：pvn3d
• Python：3.8.20
• PyTorch：1.10.0+cu113
• CUDA Toolkit：11.3
• ONNX：1.14.1
• ONNX Runtime：1.16.3
• TensorRT Python：8.6.1
• trtexec：/opt/tensorrt/TensorRT-8.6.1.6/bin/trtexec

当前环境已经满足 ONNX 导出和 ONNX 基础校验的前提。

3. 当前导出目标

当前 deploy/scripts/export_onnx.py 只导出两个子图：

1. rgb_backbone
2. fusion_head

PointNet2 不在 ONNX 导出范围内，原因是它依赖自定义 CUDA 扩展。

3.1 为什么不是把原始 graph 一次性导成一个 ONNX

当前 PVN3D 的推理路径不是一个“纯标准算子”的网络，而是混合结构：

1. RGB 主干网络
2. PointNet2 点云分支
3. RGB 与点云融合头

这里真正的问题不在 CNN，也不在最终的融合头，而在 PointNet2。

PointNet2 当前依赖本仓库里的自定义 CUDA 扩展，包括这类算子：

• furthest_point_sample
• gather_points
• ball_query
• group_points
• three_nn
• three_interpolate

这些算子当前是通过本地编译的扩展模块执行的，不是标准 ONNX 算子。

因此如果直接把整个 PVN3D 原始 graph 一次性导出，会立刻遇到两个工程问题：

1. PyTorch 到 ONNX 这一层无法稳定表达 PointNet2 的自定义 CUDA 算子
2. 即使勉强导出，后面的 TensorRT 也不能直接解析这些自定义算子，除非单独开发 TensorRT plugin

这就是当前部署路线不把整网直接导成单个 ONNX 的原因。

3.2 为什么最终变成“2 个 ONNX + 1 段 CUDA”

当前拆分后的部署结构是：

1. rgb_backbone.onnx
2. PointNet2 Native CUDA
3. fusion_head.onnx

也就是：

• 图像分支导成 ONNX
• 点云分支继续保留原生 CUDA
• 融合与预测头再导成 ONNX

这样拆的目的很明确：

• 把能稳定导出的部分先导出来
• 把最难处理的 PointNet2 自定义算子留在现有 CUDA 路径
• 先跑通第一阶段部署，再决定是否值得做 TensorRT plugin

换句话说，当前不是“原始 graph 被随意拆开”，而是按算子可导出性做的工程拆分。

3.3 为什么 PointNet2 继续保留 CUDA，而不是也转成 ONNX

原因不是 PointNet2 不重要，恰恰相反，是因为它最特殊。

当前 PointNet2 这一路已经在容器里验证通过：

• 扩展模块能导入
• CUDA 可用
• 原生推理链可工作

这里有一个验证细节需要单独说明：

• 当前容器里验证 _ext 时，应先 import torch
• 如果直接裸导入 _ext，可能报 ImportError: libc10.so: cannot open shared object file
• 这属于 PyTorch 共享库的动态加载顺序问题，不表示扩展没有编译成功

在这种前提下，继续保留 PointNet2 的原生 CUDA 路径，收益很直接：

• 不需要先做 ONNX 自定义算子适配
• 不需要先做 TensorRT plugin
• 可以先让部署链往前走

这条路线的代价也很明确：

• 最终不是单一 engine
• 推理链会变成混合执行

但对当前阶段来说，这是合理代价，因为它显著降低了第一版部署风险。

相关脚本：

• export_onnx.py
• model_wrappers.py
• run_linemod_hybrid_ort.py

4. 进入环境

docker exec -it pvn3d-dev bash
source /opt/conda/etc/profile.d/conda.sh
conda activate pvn3d
cd /workspace/workflow/self/PVN3D

建议先做一轮快速检查：

python --version
python - <<'PY'
import torch
import onnx
import onnxruntime as ort
print("torch:", torch.__version__)
print("torch_cuda:", torch.version.cuda)
print("cuda_available:", torch.cuda.is_available())
print("onnx:", onnx.__version__)
print("onnxruntime:", ort.__version__)
PY

5. ONNX 导出脚本说明

导出脚本在：

• export_onnx.py

当前关键参数：

• --checkpoint
• --output-dir
• --target
• --num-classes
• --num-points
• --height
• --width
• --device
• --opset

当前脚本已经调整过默认 opset：

• 默认值：13

原因不是偏好问题，而是当前容器里的 torch 1.10.0 不支持 opset 17 导出。

6. 推荐导出命令

当前已经验证通过的一组导出尺寸是：

• height=480
• width=624

下面的命令统一按这组尺寸记录。
之前的 480x640 是问题定位阶段的测试尺寸，不再作为当前推荐命令。

6.0 用于 LINEMOD ape 混合测试的推荐导出命令

如果后续要用 ONNX Runtime 测试 LINEMOD ape 的整条混合推理链，建议先导出与 ape 权重匹配的 ONNX：

python deploy/scripts/export_onnx.py \
  --checkpoint weights/ape_pvn3d_best.pth.tar \
  --output-dir deploy/models/onnx_ape \
  --target all \
  --num-classes 2 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --device cuda \
  --opset 13

这里不要继续使用 ycb_pvn3d_best.pth.tar 导出的 ONNX 去测 LINEMOD ape。

原因：

• YCB 是 22 类
• LINEMOD 单类配置是 2 类
• checkpoint 和输出头维度必须一致

6.1 先只测 rgb_backbone

python deploy/scripts/export_onnx.py \
  --checkpoint weights/ycb_pvn3d_best.pth.tar \
  --output-dir deploy/models/onnx \
  --target rgb_backbone \
  --num-classes 22 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --device cuda \
  --opset 13

6.2 只测 fusion_head

python deploy/scripts/export_onnx.py \
  --checkpoint weights/ycb_pvn3d_best.pth.tar \
  --output-dir deploy/models/onnx \
  --target fusion_head \
  --num-classes 22 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --device cuda \
  --opset 13

6.3 全量导出

python deploy/scripts/export_onnx.py \
  --checkpoint weights/ycb_pvn3d_best.pth.tar \
  --output-dir deploy/models/onnx \
  --target all \
  --num-classes 22 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --device cuda \
  --opset 13

建议不要一开始就跑 all。
先拆开测试，定位更快。

7. 导出后检查

如果导出成功，先检查输出目录：

ls -lh deploy/models/onnx
cat deploy/models/onnx/export_manifest.json

预期文件：

• rgb_backbone.onnx
• fusion_head.onnx
• export_manifest.json

8. ONNX 基础测试

8.1 检查模型能否被 ONNX 解析

python - <<'PY'
import onnx
model = onnx.load("deploy/models/onnx/rgb_backbone.onnx")
onnx.checker.check_model(model)
print("rgb_backbone.onnx ok")
PY

如果是 fusion_head.onnx，把文件名替换掉即可。

8.2 用 ONNX Runtime 测试能否创建会话

python - <<'PY'
import onnxruntime as ort
sess = ort.InferenceSession(
    "deploy/models/onnx/rgb_backbone.onnx",
    providers=["CPUExecutionProvider"],
)
print("inputs:", [i.name for i in sess.get_inputs()])
print("outputs:", [o.name for o in sess.get_outputs()])
PY

这里先用 CPUExecutionProvider。

原因：

• 当前目的只是检查 ONNX 图是否可加载
• 不把 ORT GPU 执行问题混进导出验证阶段

9. 混合推理全流程测试

9.1 当前是否可以用 ONNX Runtime 测整条 PVN3D 流程

可以，但要明确是“混合推理”。

当前能跑通的路径是：

1. rgb_backbone.onnx 由 ONNX Runtime 执行
2. PointNet2 由原生 PyTorch CUDA 扩展执行
3. fusion_head.onnx 由 ONNX Runtime 执行
4. 最后的 LINEMOD pose 求解仍走现有 PyTorch / numpy 逻辑

这不是整网纯 ORT，但已经覆盖了当前部署拆分后的主流程。

9.2 测试脚本位置

脚本放在：

• run_linemod_hybrid_ort.py

脚本做的事情：

1. 读取一帧 LINEMOD 样本
2. 把 RGB 从 640 裁成 624
3. 把点数从数据集原始的 12288 重采样到 4096
4. 用 ONNX Runtime 跑 rgb_backbone
5. 用原生 PointNet2 CUDA 跑点云特征
6. 用 ONNX Runtime 跑 fusion_head
7. 可选对比整模型 PyTorch 输出
8. 调用 LINEMOD pose solver，输出 pose、ADD、ADD-S

9.3 运行命令

先保证你已经导出了 ape 对应的 ONNX：

python deploy/scripts/export_onnx.py \
  --checkpoint weights/ape_pvn3d_best.pth.tar \
  --output-dir deploy/models/onnx_ape \
  --target all \
  --num-classes 2 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --device cuda \
  --opset 13

然后执行测试脚本：

python onnx_test/run_linemod_hybrid_ort.py \
  --cls ape \
  --checkpoint weights/ape_pvn3d_best.pth.tar \
  --rgb-onnx deploy/models/onnx_ape/rgb_backbone.onnx \
  --fusion-onnx deploy/models/onnx_ape/fusion_head.onnx \
  --sample-index 0 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --crop-left 8 \
  --output onnx_test/outputs/linemod_ape_hybrid_summary.json

9.4 输出内容

脚本会输出并保存一份 JSON，总结以下信息：

• 使用的 checkpoint
• 使用的 ONNX 文件
• ONNX Runtime providers
• 输入 shape
• 点云 shape
• 预测 pose
• ADD
• ADD-S
• 如果未加 --skip-torch-compare
  还会给出混合推理和整模型 PyTorch 结果的误差统计

9.5 为什么脚本里还要裁剪和重采样

这是当前导出配置决定的，不是额外的设计复杂度。

原因有两个：

1. 当前已验证通过的 ONNX 导出尺寸是 480x624
2. LINEMOD 数据集当前默认采样点数是 12288，而导出使用的是 4096

所以脚本里必须做两步对齐：

• RGB: 640 -> 624
• points: 12288 -> 4096

9.6 为什么需要 crop_left

crop_left 不是多余参数，它是用来定义 RGB 裁剪窗口起点的。

当前背景是：

• LINEMOD 原始 RGB 宽度是 640
• 当前导出的 ONNX 输入宽度是 624

所以测试脚本必须先把原图裁成 624 宽，才能送进 rgb_backbone.onnx。

但这里不能只裁 RGB，不处理 choose。

原因：

• choose 记录的是采样点在原始图像上的像素索引
• 一旦 RGB 被裁剪，像素坐标系就变了
• 脚本必须知道裁剪窗口是从哪一列开始的，才能把 choose 从原始 640 宽图正确映射到裁剪后的 624 宽图

这就是 crop_left 的作用：

1. 定义 RGB 裁剪起点
2. 让 choose 的索引同步完成重映射

当前默认值是：

--crop-left 8

原因：

• 原始宽度 640
• 导出宽度 624
• 两者差值是 16
• 左边裁 8、右边等价裁 8

也就是居中裁剪。

如果以后导出宽度改了，crop_left 也要跟着一起调整。
它不是固定常量，而是和导出分辨率绑定的参数。

10. 当前已经确认的问题

10.1 问题一：默认 opset 17 导致导出直接失败

实际报错：

ValueError: Unsupported ONNX opset version: 17

原因：

• 原脚本默认把 --opset 设成了 17
• 当前容器里的 torch 1.10.0+cu113 不支持这个导出版本

实际检查结果：

• stable opset: 7 到 13
• main opset: 14

解决：

• 当前环境统一使用 --opset 13
• 脚本默认值已经改成 13

相关代码位置：

• export_onnx.py:81
• export_onnx.py:42

10.2 问题二：torch.onnx.symbolic_helper 访问方式在 torch 1.10 下不稳

在把导出脚本默认 opset 改为 13 之后，又遇到过一轮脚本自身报错：

AttributeError: module 'torch.onnx' has no attribute 'symbolic_helper'

原因：

• 在当前 torch 1.10.0+cu113 环境里
• import torch.onnx.symbolic_helper as sh 是可以成功的
• 但 torch.onnx.symbolic_helper 不一定会自动挂成 torch.onnx 的属性

也就是说，这两种写法在当前环境里不是等价的：

错误写法：

torch.onnx.symbolic_helper

正确写法：

import torch.onnx.symbolic_helper as onnx_symbolic_helper

解决：

• 已经把 export_onnx.py 改成显式导入 torch.onnx.symbolic_helper
• 后续 opset 检查统一使用局部模块变量，不再通过 torch.onnx.symbolic_helper 访问

这属于脚本兼容性问题，不是模型结构问题。

10.3 问题三：rgb_backbone 导出卡在 adaptive_avg_pool2d

实际报错：

RuntimeError: Unsupported: ONNX export of operator adaptive_avg_pool2d,
since output size is not factor of input size

这是当前真正的 ONNX 导出阻塞点。

原因来自 PSPNet 的 PSP 模块。

相关代码：

• pspnet.py:20

这里使用了：

• AdaptiveAvgPool2d(output_size=(1,1))
• AdaptiveAvgPool2d(output_size=(2,2))
• AdaptiveAvgPool2d(output_size=(3,3))
• AdaptiveAvgPool2d(output_size=(6,6))

前一轮失败时使用的测试输入是：

• height=480
• width=640

在 torch 1.10 的 ONNX 导出器里，如果特征图尺寸相对于这些目标池化尺寸不是整倍数，导出就会失败。

这不是 ONNX 包没装好，也不是 TensorRT 的问题，而是 PyTorch 导出器本身的限制。

10.4 问题四：PointNet2 扩展编译成功，但直接导入 _ext 报 libc10.so

当前容器里还确认过一个容易误判的问题。

现象：

1. python setup.py build_ext --inplace 已经成功
2. _ext.cpython-38-*.so 已经生成
3. 直接执行下面的验证代码时失败：

python - <<'PY'
from pvn3d.lib.pointnet2_utils import _ext
print(_ext)
PY

报错：

ImportError: libc10.so: cannot open shared object file: No such file or directory

原因：

• _ext.so 依赖 PyTorch 的共享库
• 例如 libc10.so
• 这些库在当前 conda 环境的 torch/lib 目录下
• 如果当前 Python 进程还没有先加载 torch，直接导入 _ext 时，动态链接器可能找不到这些依赖

这说明问题不在编译结果本身，而在动态库加载顺序。

解决：

python - <<'PY'
import torch
from pvn3d.lib.pointnet2_utils import _ext
print(_ext)
PY

这个顺序已经在 pvn3d-dev 容器里验证通过。

结论：

• build_ext --inplace 成功，说明扩展已经编译出来
• 如果裸导入 _ext 报 libc10.so，优先先 import torch
• 这不是 ONNX 导出 bug，而是当前环境下 PointNet2 扩展验证时的依赖加载问题

10.5 问题五：Invalid magic number 实际上是 checkpoint 保存格式不同

在用 LINEMOD ape 权重导出 ONNX 时，出现过下面的报错：

RuntimeError: Invalid magic number; corrupt file?

对应命令：

python deploy/scripts/export_onnx.py \
  --checkpoint weights/ape_pvn3d_best.pth.tar \
  --output-dir deploy/models/onnx_ape \
  --target all \
  --num-classes 2 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --device cuda \
  --opset 13

这个报错看起来像文件损坏，但当前容器里实际确认的原因不是损坏，而是保存格式不同：

• ycb_pvn3d_best.pth.tar 可以直接用 torch.load
• ape_pvn3d_best.pth.tar 是普通 pickle checkpoint
• 其中仍然包含 model_state

所以问题不在权重内容，而在加载器一开始只支持 torch.load。

解决：

• 已经更新 model_wrappers.py
• 现在会先尝试 torch.load
• 如果遇到 Invalid magic number
• 自动回退到 pickle.load

结论：

• Invalid magic number 在这里不等于 checkpoint 损坏
• 它表示当前 checkpoint 不是 torch.save 格式
• 当前导出脚本已经兼容这两种 checkpoint

10.6 问题六：混合测试脚本找不到 obj_01.ply

在运行下面的混合测试命令时：

python onnx_test/run_linemod_hybrid_ort.py \
  --cls ape \
  --checkpoint weights/ape_pvn3d_best.pth.tar \
  --rgb-onnx deploy/models/onnx_ape/rgb_backbone.onnx \
  --fusion-onnx deploy/models/onnx_ape/fusion_head.onnx \
  --sample-index 0 \
  --num-points 4096 \
  --height 480 \
  --width 624 \
  --crop-left 8 \
  --output onnx_test/outputs/linemod_ape_hybrid_summary.json

出现过下面的报错：

FileNotFoundError: [Errno 2] No such file or directory:
'datasets/linemod/Linemod_preprocessed/models/obj_01.ply'

原因不是模型文件不存在。
当前容器里实际目录是存在的：

• /workspace/workflow/self/PVN3D/pvn3d/datasets/linemod/Linemod_preprocessed/models/obj_01.ply

真正原因是：

• basic_utils.py 里某些 LINEMOD 模型路径仍然写成了相对路径
• 相对路径是按 pvn3d/ 目录为当前工作目录来解析的
• 但 onnx_test/run_linemod_hybrid_ort.py 一开始是在仓库根目录执行

也就是说，这是 legacy 工具代码里的路径假设问题，不是数据集缺失。

解决：

• 已经更新 run_linemod_hybrid_ort.py
• 脚本启动后会自动切换当前工作目录到 pvn3d/
• 这样 datasets/linemod/... 这类相对路径就能正确解析

结论：

• 报 obj_01.ply 找不到时，不要先怀疑数据集没准备好
• 先检查是不是工作目录不符合 basic_utils.py 的相对路径假设
• 当前测试脚本已经做了自动修正

11. 当前问题的处理建议

11.1 先固定 opset 13

这是当前环境下的硬约束，不要再尝试 17。

11.2 确保使用更新后的导出脚本

当前 export_onnx.py 已经做了两处修正：

1. 默认 opset 从 17 改为 13
2. symbolic_helper 改成显式导入方式

如果你本地容器里运行的还是旧版本脚本，先同步这份脚本再测试。

11.3 PointNet2 扩展验证时先导入 torch

如果需要验证 PointNet2 扩展是否可用，使用下面的顺序：

python - <<'PY'
import torch
from pvn3d.lib.pointnet2_utils import _ext
print(_ext)
PY

不要把下面这种“裸导入 _ext”当成唯一验证方式：

python - <<'PY'
from pvn3d.lib.pointnet2_utils import _ext
print(_ext)
PY

在当前容器里，这种写法可能因为 libc10.so 的动态加载顺序而失败。

11.4 使用更新后的 checkpoint 加载逻辑

如果导出 LINEMOD ape 权重时遇到：

Invalid magic number; corrupt file?

先不要判断为权重损坏。
当前应先确认使用的是更新后的：

• model_wrappers.py

现在它已经兼容：

• torch.save checkpoint
• pickle checkpoint

11.5 混合测试脚本应使用更新后的版本

如果运行 onnx_test/run_linemod_hybrid_ort.py 时出现：

FileNotFoundError: ... obj_01.ply

先确认使用的是更新后的脚本版本。
当前脚本已经自动切到 pvn3d/ 工作目录，避免 basic_utils.py 的相对路径解析失败。

11.6 优先尝试更适合 PSP 导出的固定分辨率

当前已经验证通过的尺寸是：

--height 480 --width 624

原因：

• PSP 模块的池化目标是 1, 2, 3, 6
• 这类结构更适合特征图宽高都能被 1,2,3,6 整除的输入
• 480 和 624 这一组尺寸已经在当前容器里完成过 ONNX 导出
• 相比 480x640，这一组更容易满足 PSP pooling 的导出约束

对当前容器，这一步不再是“尝试”，而是目前已经跑通的推荐值。

11.7 如果尺寸不能改，就做导出专用改写

如果部署输入必须固定为 480x640，那就不要继续硬试默认 PSP 实现了。

更直接的做法是：

• 为 ONNX 导出单独改写 PSPModule
• 把 AdaptiveAvgPool2d 换成更容易被 torch 1.10 导出的显式 pooling 方案

这条路需要改模型代码，但从工程上比升级整套 PyTorch 环境更稳。

11.8 不建议优先升级 PyTorch

原因：

• 会影响当前稳定的 torch 1.10 + cu113
• 会影响 PointNet2 扩展
• 会把问题从“局部导出兼容性”扩大成“整套训练/推理环境迁移”

12. 当前建议的测试顺序

建议按下面顺序执行：

1. python deploy/scripts/export_onnx.py --help
2. 用 ape checkpoint 导出 deploy/models/onnx_ape
3. 只测试 fusion_head
4. 再测试 rgb_backbone
5. rgb_backbone 导出失败时，先改输入尺寸
6. 导出成功后，用 onnx.checker 检查
7. 再用 ONNX Runtime 建立会话
8. 再跑 onnx_test/run_linemod_hybrid_ort.py
9. 最后再进入 TensorRT engine 构建

这样做的原因：

• fusion_head 不包含 PSPNet 的自适应池化问题
• 先拆开验证，能更快判断问题到底在 RGB backbone 还是 fusion head

13. 当前状态结论

截至目前，ONNX 导出与测试状态可以概括成下面几句：

1. 当前可以测试整条 PVN3D 主流程，但方式是“混合推理”，不是整网纯 ORT
2. rgb_backbone 在 480x640 下主要卡在 PSP 模块的 AdaptiveAvgPool2d 导出限制
3. 在 480x624 下，当前容器已经成功导出 rgb_backbone.onnx 和 fusion_head.onnx
4. PointNet2 继续保留原生 CUDA，是当前这条测试链成立的前提
5. 当前阻塞点已经不再是 ONNX 缺包，而是后续要继续验证混合推理结果与 TensorRT engine 构建

补一条结构层面的结论：

6. 当前部署路径本来就不是“整网一个 ONNX”，而是“两个 ONNX + PointNet2 原生 CUDA”

下一步最合理的动作是：

• 用 ape 权重重新导出匹配 LINEMOD 的 ONNX
• 跑 onnx_test/run_linemod_hybrid_ort.py
• 看混合推理和整模型 PyTorch 的误差
• 再决定是否继续推进 TensorRT engine 构建