1. 目的
本文单独记录当前 pvn3d-dev 容器内的 ONNX 转换流程、测试方法,以及已经确认的导出问题。
这篇文档只关注两件事:
- 如何把当前
deploy/里的可导出子图转成 ONNX - 导出过程中实际遇到了什么问题,如何判断,下一步怎么处理
不在本文展开的内容:
- TensorRT 安装
- 容器创建
- PointNet2 插件化
这里先说明一个边界:
- 当前可以用 ONNX Runtime 测试整条 PVN3D 推理链
- 但不是"整网纯 ONNX Runtime"
- 当前可行的是"ONNX Runtime + PointNet2 Native CUDA"的混合推理流程
原因很简单:
rgb_backbone可以走 ONNX Runtimefusion_head可以走 ONNX Runtime- PointNet2 仍然保留 PyTorch CUDA 扩展
相关文档:
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 只导出两个子图:
rgb_backbonefusion_head
PointNet2 不在 ONNX 导出范围内,原因是它依赖自定义 CUDA 扩展。
3.1 为什么不是把原始 graph 一次性导成一个 ONNX
当前 PVN3D 的推理路径不是一个"纯标准算子"的网络,而是混合结构:
- RGB 主干网络
- PointNet2 点云分支
- RGB 与点云融合头
这里真正的问题不在 CNN,也不在最终的融合头,而在 PointNet2。
PointNet2 当前依赖本仓库里的自定义 CUDA 扩展,包括这类算子:
furthest_point_samplegather_pointsball_querygroup_pointsthree_nnthree_interpolate
这些算子当前是通过本地编译的扩展模块执行的,不是标准 ONNX 算子。
因此如果直接把整个 PVN3D 原始 graph 一次性导出,会立刻遇到两个工程问题:
- PyTorch 到 ONNX 这一层无法稳定表达 PointNet2 的自定义 CUDA 算子
- 即使勉强导出,后面的 TensorRT 也不能直接解析这些自定义算子,除非单独开发 TensorRT plugin
这就是当前部署路线不把整网直接导成单个 ONNX 的原因。
3.2 为什么最终变成"2 个 ONNX + 1 段 CUDA"
当前拆分后的部署结构是:
rgb_backbone.onnxPointNet2 Native CUDAfusion_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
- 推理链会变成混合执行
但对当前阶段来说,这是合理代价,因为它显著降低了第一版部署风险。
相关脚本:
4. 进入环境
bash
docker exec -it pvn3d-dev bash
source /opt/conda/etc/profile.d/conda.sh
conda activate pvn3d
cd /workspace/workflow/self/PVN3D
建议先做一轮快速检查:
bash
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 导出脚本说明
导出脚本在:
当前关键参数:
--checkpoint--output-dir--target--num-classes--num-points--height--width--device--opset
当前脚本已经调整过默认 opset:
- 默认值:
13
原因不是偏好问题,而是当前容器里的 torch 1.10.0 不支持 opset 17 导出。
6. 推荐导出命令
当前已经验证通过的一组导出尺寸是:
height=480width=624
下面的命令统一按这组尺寸记录。
之前的 480x640 是问题定位阶段的测试尺寸,不再作为当前推荐命令。
6.0 用于 LINEMOD ape 混合测试的推荐导出命令
如果后续要用 ONNX Runtime 测试 LINEMOD ape 的整条混合推理链,建议先导出与 ape 权重匹配的 ONNX:
bash
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
bash
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
bash
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 全量导出
bash
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. 导出后检查
如果导出成功,先检查输出目录:
bash
ls -lh deploy/models/onnx
cat deploy/models/onnx/export_manifest.json
预期文件:
rgb_backbone.onnxfusion_head.onnxexport_manifest.json
8. ONNX 基础测试
8.1 检查模型能否被 ONNX 解析
bash
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 测试能否创建会话
bash
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 流程
可以,但要明确是"混合推理"。
当前能跑通的路径是:
rgb_backbone.onnx由 ONNX Runtime 执行- PointNet2 由原生 PyTorch CUDA 扩展执行
fusion_head.onnx由 ONNX Runtime 执行- 最后的 LINEMOD pose 求解仍走现有 PyTorch / numpy 逻辑
这不是整网纯 ORT,但已经覆盖了当前部署拆分后的主流程。
9.2 测试脚本位置
脚本放在:
脚本做的事情:
- 读取一帧 LINEMOD 样本
- 把 RGB 从
640裁成624 - 把点数从数据集原始的
12288重采样到4096 - 用 ONNX Runtime 跑
rgb_backbone - 用原生 PointNet2 CUDA 跑点云特征
- 用 ONNX Runtime 跑
fusion_head - 可选对比整模型 PyTorch 输出
- 调用 LINEMOD pose solver,输出 pose、ADD、ADD-S
9.3 运行命令
先保证你已经导出了 ape 对应的 ONNX:
bash
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
然后执行测试脚本:
bash
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 为什么脚本里还要裁剪和重采样
这是当前导出配置决定的,不是额外的设计复杂度。
原因有两个:
- 当前已验证通过的 ONNX 导出尺寸是
480x624 - 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 的作用:
- 定义 RGB 裁剪起点
- 让
choose的索引同步完成重映射
当前默认值是:
bash
--crop-left 8
原因:
- 原始宽度
640 - 导出宽度
624 - 两者差值是
16 - 左边裁
8、右边等价裁8
也就是居中裁剪。
如果以后导出宽度改了,crop_left 也要跟着一起调整。
它不是固定常量,而是和导出分辨率绑定的参数。
10. 当前已经确认的问题
10.1 问题一:默认 opset 17 导致导出直接失败
实际报错:
text
ValueError: Unsupported ONNX opset version: 17
原因:
- 原脚本默认把
--opset设成了17 - 当前容器里的
torch 1.10.0+cu113不支持这个导出版本
实际检查结果:
- stable opset:
7到13 - main opset:
14
解决:
- 当前环境统一使用
--opset 13 - 脚本默认值已经改成
13
相关代码位置:
10.2 问题二:torch.onnx.symbolic_helper 访问方式在 torch 1.10 下不稳
在把导出脚本默认 opset 改为 13 之后,又遇到过一轮脚本自身报错:
text
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的属性
也就是说,这两种写法在当前环境里不是等价的:
错误写法:
python
torch.onnx.symbolic_helper
正确写法:
python
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
实际报错:
text
RuntimeError: Unsupported: ONNX export of operator adaptive_avg_pool2d,
since output size is not factor of input size
这是当前真正的 ONNX 导出阻塞点。
原因来自 PSPNet 的 PSP 模块。
相关代码:
这里使用了:
AdaptiveAvgPool2d(output_size=(1,1))AdaptiveAvgPool2d(output_size=(2,2))AdaptiveAvgPool2d(output_size=(3,3))AdaptiveAvgPool2d(output_size=(6,6))
前一轮失败时使用的测试输入是:
height=480width=640
在 torch 1.10 的 ONNX 导出器里,如果特征图尺寸相对于这些目标池化尺寸不是整倍数,导出就会失败。
这不是 ONNX 包没装好,也不是 TensorRT 的问题,而是 PyTorch 导出器本身的限制。
10.4 问题四:PointNet2 扩展编译成功,但直接导入 _ext 报 libc10.so
当前容器里还确认过一个容易误判的问题。
现象:
python setup.py build_ext --inplace已经成功_ext.cpython-38-*.so已经生成- 直接执行下面的验证代码时失败:
bash
python - <<'PY'
from pvn3d.lib.pointnet2_utils import _ext
print(_ext)
PY
报错:
text
ImportError: libc10.so: cannot open shared object file: No such file or directory
原因:
_ext.so依赖 PyTorch 的共享库- 例如
libc10.so - 这些库在当前 conda 环境的
torch/lib目录下 - 如果当前 Python 进程还没有先加载
torch,直接导入_ext时,动态链接器可能找不到这些依赖
这说明问题不在编译结果本身,而在动态库加载顺序。
解决:
bash
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 时,出现过下面的报错:
text
RuntimeError: Invalid magic number; corrupt file?
对应命令:
bash
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.loadape_pvn3d_best.pth.tar是普通picklecheckpoint- 其中仍然包含
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
在运行下面的混合测试命令时:
bash
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
出现过下面的报错:
text
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 已经做了两处修正:
- 默认
opset从17改为13 symbolic_helper改成显式导入方式
如果你本地容器里运行的还是旧版本脚本,先同步这份脚本再测试。
11.3 PointNet2 扩展验证时先导入 torch
如果需要验证 PointNet2 扩展是否可用,使用下面的顺序:
bash
python - <<'PY'
import torch
from pvn3d.lib.pointnet2_utils import _ext
print(_ext)
PY
不要把下面这种"裸导入 _ext"当成唯一验证方式:
bash
python - <<'PY'
from pvn3d.lib.pointnet2_utils import _ext
print(_ext)
PY
在当前容器里,这种写法可能因为 libc10.so 的动态加载顺序而失败。
11.4 使用更新后的 checkpoint 加载逻辑
如果导出 LINEMOD ape 权重时遇到:
text
Invalid magic number; corrupt file?
先不要判断为权重损坏。
当前应先确认使用的是更新后的:
现在它已经兼容:
torch.savecheckpointpicklecheckpoint
11.5 混合测试脚本应使用更新后的版本
如果运行 onnx_test/run_linemod_hybrid_ort.py 时出现:
text
FileNotFoundError: ... obj_01.ply
先确认使用的是更新后的脚本版本。
当前脚本已经自动切到 pvn3d/ 工作目录,避免 basic_utils.py 的相对路径解析失败。
11.6 优先尝试更适合 PSP 导出的固定分辨率
当前已经验证通过的尺寸是:
bash
--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. 当前建议的测试顺序
建议按下面顺序执行:
python deploy/scripts/export_onnx.py --help- 用
apecheckpoint 导出deploy/models/onnx_ape - 只测试
fusion_head - 再测试
rgb_backbone rgb_backbone导出失败时,先改输入尺寸- 导出成功后,用
onnx.checker检查 - 再用 ONNX Runtime 建立会话
- 再跑
onnx_test/run_linemod_hybrid_ort.py - 最后再进入 TensorRT engine 构建
这样做的原因:
fusion_head不包含 PSPNet 的自适应池化问题- 先拆开验证,能更快判断问题到底在 RGB backbone 还是 fusion head
13. 当前状态结论
截至目前,ONNX 导出与测试状态可以概括成下面几句:
- 当前可以测试整条 PVN3D 主流程,但方式是"混合推理",不是整网纯 ORT
rgb_backbone在480x640下主要卡在 PSP 模块的AdaptiveAvgPool2d导出限制- 在
480x624下,当前容器已经成功导出rgb_backbone.onnx和fusion_head.onnx - PointNet2 继续保留原生 CUDA,是当前这条测试链成立的前提
- 当前阻塞点已经不再是 ONNX 缺包,而是后续要继续验证混合推理结果与 TensorRT engine 构建
补一条结构层面的结论:
- 当前部署路径本来就不是"整网一个 ONNX",而是"两个 ONNX + PointNet2 原生 CUDA"
下一步最合理的动作是:
- 用
ape权重重新导出匹配 LINEMOD 的 ONNX - 跑
onnx_test/run_linemod_hybrid_ort.py - 看混合推理和整模型 PyTorch 的误差
- 再决定是否继续推进 TensorRT engine 构建