PVN3D TensorRT Engine 转换与测试记录

1. 目的

本文单独记录当前 pvn3d-dev 容器内，从 ONNX 到 TensorRT engine 的转换流程、测试方法，以及这一阶段需要注意的约束。

这篇文档只关注两件事：

如何把当前已经导出的 ONNX 子图转换成 TensorRT engine
如何验证 engine 是否构建成功，以及这条链路当前有哪些边界

不在本文展开的内容：

TensorRT 安装
ONNX 导出问题本身
PointNet2 TensorRT plugin

2. 当前容器环境

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

容器：pvn3d-dev
仓库路径：/workspace/workflow/self/PVN3D
Conda 环境：pvn3d
Python：3.8.20
PyTorch：1.10.0+cu113
ONNX：1.14.1
ONNX Runtime：1.16.3
TensorRT Python：8.6.1
trtexec：/opt/tensorrt/TensorRT-8.6.1.6/bin/trtexec
trtexec --help：可执行，版本标识为 TensorRT v8601

当前环境已经满足 engine 构建前提。

3. 当前 TensorRT 转换目标

当前不是把整个 PVN3D 一次性转换成单个 TensorRT engine。

当前转换目标只有两个：

rgb_backbone.onnx -> rgb_backbone.engine
fusion_head.onnx -> fusion_head.engine

PointNet2 仍然保留原生 CUDA 路径。

所以当前实际部署结构仍然是：

rgb_backbone.engine
PointNet2 Native CUDA
fusion_head.engine

这条路线延续了 ONNX 拆分阶段的同一原则：

先把标准算子图转换成 TensorRT，把 PointNet2 自定义 CUDA 算子留在现有路径。

4. 当前已有 ONNX 产物

当前仓库里已经存在两组 ONNX：

deploy/models/onnx/*.onnx
deploy/models/onnx_ape/*.onnx

对 LINEMOD ape 的测试，应优先使用：

deploy/models/onnx_ape/rgb_backbone.onnx
deploy/models/onnx_ape/fusion_head.onnx

原因：

onnx_ape 是用 ape_pvn3d_best.pth.tar 导出的
num_classes=2
与 LINEMOD ape 的单类配置一致

不要拿 YCB 22 类 的 ONNX 去构建 LINEMOD ape 的 engine。

这里也补一条边界说明：

当前 LINEMOD 单物体流程里，num_classes 通常固定为 2
所以如果后续从 ape 换到其他 LINEMOD 物体，通常不需要修改 num_classes

例如：

ape：num_classes=2
can：num_classes=2
driller：num_classes=2

但即便如此，下面这些仍然不能混用：

checkpoint
ONNX 目录
engine 目录
测试脚本里的 --cls

原因不是类别数不同，而是虽然都属于"单前景类"，这个前景类对应的具体物体不同：

权重不同
mesh 不同
keypoints 不同
center 不同
pose 求解目标也不同

所以对 LINEMOD 更准确的理解是：

num_classes=2 往往不变
但 ape、can、driller 仍应各自导出各自的 ONNX 和 engine

5. 进入环境

bash 复制代码

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

建议先确认：

bash 复制代码

which trtexec
python - <<'PY'
import tensorrt as trt
print("TensorRT:", trt.__version__)
PY

6. 当前 engine 构建脚本

当前仓库使用：

build_trt_engine.py

这个脚本本质上是对 trtexec 的封装。

它当前负责：

检查 trtexec 是否存在
在需要时按目标子图自动拼接输入 shape
生成并执行 trtexec 命令
输出 .engine

当前支持目标：

rgb_backbone
fusion_head

7. 构建前的关键约束

7.0 先区分静态 ONNX 和动态 ONNX

当前 deploy/models/onnx_ape/*.onnx 已经确认是静态输入模型：

rgb_backbone.onnx
- rgb: [1, 3, 480, 624]
fusion_head.onnx
- out_rgb: [1, 128, 480, 624]
- choose: [1, 1, 4096]
- pcld_emb: [1, 128, 4096]

这意味着当前构建 engine 时，不应该再给 trtexec 传：

--minShapes
--optShapes
--maxShapes

如果对静态 ONNX 强行传这些参数，TensorRT 会直接报错。

7.1 分辨率必须与导出 ONNX 一致

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

height=480
width=624

所以构建 TensorRT engine 时，也必须使用同样的输入尺寸。

不要再写成：

bash 复制代码

--height 480 --width 640

当前推荐值是：

bash 复制代码

--height 480 --width 624

7.2 点数必须与导出 ONNX 一致

当前导出使用的是：

bash 复制代码

--num-points 4096

所以 fusion_head.engine 构建时，也必须保持：

bash 复制代码

--num-points 4096

7.3 ONNX 和 engine 的类别语义必须一致

如果构建 LINEMOD ape 的 engine，优先使用：

deploy/models/onnx_ape

不要把 deploy/models/onnx 和 deploy/models/onnx_ape 混着用。

8. 推荐构建命令

8.1 先做 dry-run

先只打印命令，不实际构建：

bash 复制代码

python deploy/scripts/build_trt_engine.py \
  --onnx-dir deploy/models/onnx_ape \
  --engine-dir deploy/models/engine_ape \
  --target all \
  --fp16 \
  --height 480 \
  --width 624 \
  --num-points 4096 \
  --dry-run

这样做的目的：

先确认 onnx-dir 是否正确
先确认静态 ONNX 没有被错误地附加动态 profile shape
先确认脚本实际拼出来的 trtexec 命令是否符合预期

8.2 实际构建

确认 dry-run 没问题后，再执行：

bash 复制代码

python deploy/scripts/build_trt_engine.py \
  --onnx-dir deploy/models/onnx_ape \
  --engine-dir deploy/models/engine_ape \
  --target all \
  --fp16 \
  --height 480 \
  --width 624 \
  --num-points 4096

8.3 只构建 `rgb_backbone`

bash 复制代码

python deploy/scripts/build_trt_engine.py \
  --onnx-dir deploy/models/onnx_ape \
  --engine-dir deploy/models/engine_ape \
  --target rgb_backbone \
  --fp16 \
  --height 480 \
  --width 624 \
  --num-points 4096

8.4 只构建 `fusion_head`

bash 复制代码

python deploy/scripts/build_trt_engine.py \
  --onnx-dir deploy/models/onnx_ape \
  --engine-dir deploy/models/engine_ape \
  --target fusion_head \
  --fp16 \
  --height 480 \
  --width 624 \
  --num-points 4096

9. 当前脚本对 shape 的处理逻辑

9.1 静态 ONNX

如果 ONNX 输入维度都是固定值，脚本现在会：

不传 --minShapes
不传 --optShapes
不传 --maxShapes

这是当前 onnx_ape 的实际情况。

例如 rgb_backbone 的 dry-run 结果现在是：

bash 复制代码

trtexec --onnx=deploy/models/onnx_ape/rgb_backbone.onnx \
  --saveEngine=deploy/models/engine_ape/rgb_backbone.engine \
  --memPoolSize=workspace:4096 \
  --skipInference \
  --fp16

9.2 动态 ONNX

只有当 ONNX 本身包含动态维度时，脚本才会补：

--minShapes
--optShapes
--maxShapes

当前这套 onnx_ape 不属于这种情况。

9.3 如果模型是动态 ONNX，目标 shape 约定如下

脚本会传给 trtexec：

text 复制代码

rgb:1x3x480x624

对应：

input: rgb
shape: [1, 3, 480, 624]

9.4 `fusion_head`

脚本会传给 trtexec：

text 复制代码

out_rgb:1x128x480x624
choose:1x1x4096
pcld_emb:1x128x4096

也就是：

out_rgb: [1, 128, 480, 624]
choose: [1, 1, 4096]
pcld_emb: [1, 128, 4096]

这与当前导出脚本的输入约定保持一致。

10. 构建后检查

构建成功后，先检查输出目录：

bash 复制代码

find deploy/models/engine_ape -maxdepth 1 -type f | sort

预期文件：

deploy/models/engine_ape/rgb_backbone.engine
deploy/models/engine_ape/fusion_head.engine

11. engine 基础测试

11.1 用 `trtexec` 加载 engine

先单独测 rgb_backbone.engine：

bash 复制代码

trtexec \
  --loadEngine=deploy/models/engine_ape/rgb_backbone.engine \
  --skipInference

然后测 fusion_head.engine：

bash 复制代码

trtexec \
  --loadEngine=deploy/models/engine_ape/fusion_head.engine \
  --skipInference

这里的目的不是测性能，而是确认：

engine 能被 TensorRT 正常反序列化
没有明显的版本或 shape 错误

11.2 看帮助确认支持项

当前容器里的 trtexec --help 已经验证通过，说明 CLI 可用。

如果后续要补性能测试，可以继续使用：

--dumpProfile
--dumpLayerInfo
--shapes=...

12. 与混合推理测试的关系

这篇文档只负责"ONNX -> engine"的构建与最小验证。

如果你要验证完整推理链，应结合：

当前更合理的推进顺序是：

ONNX 导出成功
ONNX Runtime 能加载 ONNX
trtexec 能成功构建 engine
trtexec --loadEngine 能成功加载 engine
再把 ONNX Runtime 版混合推理逐步替换成 TensorRT engine 版混合推理

13. TensorRT Python 混合测试

13.1 当前是否可以写 TensorRT Python 测试脚本

可以。

当前容器里已经确认：

tensorrt==8.6.1
torch==1.10.0+cu113
没有 pycuda
没有 cupy

因此当前 Python 测试脚本采用的是：

TensorRT Python API 负责反序列化 engine 和执行
torch.cuda 负责分配 GPU tensor
通过 tensor.data_ptr() 把地址传给 TensorRT

这条方式不依赖 pycuda。

13.2 脚本位置

脚本放在：

run_linemod_hybrid_trt.py

脚本做的事情：

读取一帧 LINEMOD 样本
把 RGB 从 640 裁成 624
把点数从 12288 重采样到 4096
用 TensorRT engine 跑 rgb_backbone
用原生 PointNet2 CUDA 跑点云特征
用 TensorRT engine 跑 fusion_head
可选对比整模型 PyTorch 输出
调用 LINEMOD pose solver，输出 pose、ADD、ADD-S

13.3 运行前提

这个脚本依赖两个 engine 都已经存在：

deploy/models/engine_ape/rgb_backbone.engine
deploy/models/engine_ape/fusion_head.engine

少一个都不成立。

13.4 运行命令

bash 复制代码

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

这里可以继续使用相对路径参数，不需要手工改成绝对路径。

当前脚本已经会在启动时自动把下面这些参数解析成仓库根目录下的绝对路径：

--checkpoint
--rgb-engine
--fusion-engine
--output

这样做的原因是：

脚本内部为了兼容 LINEMOD 的 legacy 工具代码，会切工作目录到 pvn3d/
如果不先把这些路径转成绝对路径，weights/...、deploy/... 这类相对路径会被带偏

所以当前推荐用法仍然是文档里的相对路径写法。

13.5 输出内容

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

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

14. 当前已确认的问题

14.1 `build_trt_engine.py` 的默认宽度还是 `640`

当前脚本里的默认参数仍是：

python 复制代码

parser.add_argument("--width", type=int, default=640, help="图像宽。")

这和当前已经验证通过的 ONNX 尺寸 480x624 不一致。

这不表示脚本不能用，但表示：

运行时必须显式传 --width 624
不能依赖默认值

14.2 静态 ONNX 不能传显式 shape profile

当前已经实际遇到过下面的失败：

text 复制代码

Static model does not take explicit shapes since the shape of inference tensors
will be determined by the model itself

触发原因：

rgb_backbone.onnx 是静态模型
但构建时仍传了 --minShapes/--optShapes/--maxShapes

这不是 ONNX 损坏，也不是 TensorRT 安装有问题，而是 trtexec 参数和模型类型不匹配。

解决：

已经更新 build_trt_engine.py
现在脚本会自动检查 ONNX 是否是静态输入
对静态 ONNX，不再传动态 profile shape

14.3 `trtexec --version` 不会只打印版本

在当前 TensorRT 8.6.1 环境里：

bash 复制代码

trtexec --version

会先打印版本标识：

text 复制代码

TensorRT v8601

但后面仍可能继续进入参数解析，并报缺少模型。

这不表示 TensorRT 有问题，只是这个 CLI 的行为方式如此。

14.4 TRT Python 测试脚本的路径也会受 `chdir` 影响

当前 run_linemod_hybrid_trt.py 内部会切到 pvn3d/ 工作目录，以兼容 LINEMOD 旧工具代码里的相对路径假设。

如果脚本不先处理外部传入路径，就会出现这类问题：

weights/ape_pvn3d_best.pth.tar 被错误解析
deploy/models/engine_ape/... 被错误解析

当前脚本已经修复：

先把外部传入路径转成仓库根目录下的绝对路径
再切工作目录

因此现在文档里的相对路径命令是可直接使用的。

15. 当前处理建议

15.1 构建 engine 时显式传入 `480x624`

不要依赖脚本默认值。

15.2 对 LINEMOD `ape` 使用 `onnx_ape`

不要混用 YCB 的 ONNX。

15.3 先 dry-run 再实际构建

这样最容易发现 shape 和路径错误。

15.4 对静态 ONNX，不要手工再补动态 profile shape

如果脚本已经识别为静态 ONNX，就不要手工往命令里加：

--minShapes
--optShapes
--maxShapes

15.5 先做最小 engine 验证，再做整链替换

当前阶段优先确认：

ONNX 正常
engine 正常
TensorRT 反序列化正常

然后再继续做混合推理替换。

15.6 有两个 engine 以后再跑 TRT Python 混合测试

当前 TRT Python 测试脚本依赖：

rgb_backbone.engine
fusion_head.engine

少一个都不成立。

15.7 直接按文档命令运行，不需要手工改绝对路径

当前文档里的 TRT Python 测试命令已经与脚本实现保持一致。

继续使用：

weights/...
deploy/models/...
onnx_test/outputs/...

这种相对路径写法即可。

16. 当前建议的测试顺序

建议按下面顺序执行：

确认 deploy/models/onnx_ape/*.onnx 存在
python deploy/scripts/build_trt_engine.py ... --dry-run
python deploy/scripts/build_trt_engine.py ...
trtexec --loadEngine=... --skipInference
python onnx_test/run_linemod_hybrid_trt.py ...
再考虑把 ORT 混合推理替换成 TRT 混合推理

PVN3D TensorRT Engine 转换与测试记录

1. 目的

2. 当前容器环境

3. 当前 TensorRT 转换目标

4. 当前已有 ONNX 产物

5. 进入环境

6. 当前 engine 构建脚本

7. 构建前的关键约束

7.0 先区分静态 ONNX 和动态 ONNX

7.1 分辨率必须与导出 ONNX 一致

7.2 点数必须与导出 ONNX 一致

7.3 ONNX 和 engine 的类别语义必须一致

8. 推荐构建命令

8.1 先做 dry-run

8.2 实际构建

8.3 只构建 rgb_backbone

8.4 只构建 fusion_head

9. 当前脚本对 shape 的处理逻辑

9.1 静态 ONNX

9.2 动态 ONNX

9.3 如果模型是动态 ONNX，目标 shape 约定如下

9.4 fusion_head

10. 构建后检查

11. engine 基础测试

11.1 用 trtexec 加载 engine

11.2 看帮助确认支持项

12. 与混合推理测试的关系

13. TensorRT Python 混合测试

13.1 当前是否可以写 TensorRT Python 测试脚本

13.2 脚本位置

13.3 运行前提

13.4 运行命令

13.5 输出内容

14. 当前已确认的问题

14.1 build_trt_engine.py 的默认宽度还是 640

14.2 静态 ONNX 不能传显式 shape profile

14.3 trtexec --version 不会只打印版本

14.4 TRT Python 测试脚本的路径也会受 chdir 影响

15. 当前处理建议

15.1 构建 engine 时显式传入 480x624

15.2 对 LINEMOD ape 使用 onnx_ape

15.3 先 dry-run 再实际构建

15.4 对静态 ONNX，不要手工再补动态 profile shape

15.5 先做最小 engine 验证，再做整链替换

15.6 有两个 engine 以后再跑 TRT Python 混合测试

15.7 直接按文档命令运行，不需要手工改绝对路径

16. 当前建议的测试顺序

8.3 只构建 `rgb_backbone`

8.4 只构建 `fusion_head`

9.4 `fusion_head`

11.1 用 `trtexec` 加载 engine

14.1 `build_trt_engine.py` 的默认宽度还是 `640`

14.3 `trtexec --version` 不会只打印版本

14.4 TRT Python 测试脚本的路径也会受 `chdir` 影响

15.1 构建 engine 时显式传入 `480x624`

15.2 对 LINEMOD `ape` 使用 `onnx_ape`