本文档记录在 Windows 11 + RTX 5060 Ti (Blackwell SM_120a) 环境下,将 SAM3 模型部署为 TensorRT GPU 推理引擎的完整过程,以及 FP16 精度问题的排查结果。
一、环境安装
1. 驱动与 CUDA
| 组件 | 版本 | 下载 |
|---|---|---|
| GPU 驱动 | 591.86+ | NVIDIA 官网 |
| CUDA Toolkit | 12.8 | https://developer.nvidia.com/cuda-downloads |
CUDA 12.8 安装后路径为:
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8
将以下路径加入系统 PATH:
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\lib\x64
2. cuDNN
cuDNN 按 CUDA 大版本打包(如 cuDNN 9.x for CUDA 12.x),不分小版本,所以 CUDA 12.9 的 cuDNN 可以直接用于 CUDA 12.8。
安装方式 :下载对应 CUDA 12.x 的 cuDNN zip,将 bin/、lib/、include/ 文件覆盖到 CUDA 安装目录。
当前环境使用:
- cuDNN 8.9.7 (CUDA 12.8 自带,位于
bin/cudnn64_8.dll) - cuDNN 9.23.0 (从 CUDA 12.9 包覆盖,新增
cudnn64_9.dll系列)
两个版本共存,TensorRT 实际调用哪个取决于它链接的是 cudnn64_8.dll 还是 cudnn64_9.dll。
3. TensorRT
项目使用两个版本对比:
| 版本 | 路径 | 用途 |
|---|---|---|
| TRT-RTX 1.5.0.114 | C:\My_Project\src\deeplearning\SAM3\TensorRT-RTX-1.5.0.114\ |
主要部署版本 |
| TRT 10.11.0.33 | C:\package\TensorRT-10.11.0.33\ |
对比测试 |
Windows 环境配置(Python 脚本中):
python
import os
# TRT-RTX 路径
os.environ["PATH"] = r"C:\My_Project\src\deeplearning\SAM3\TensorRT-RTX-1.5.0.114\bin;" + os.environ.get("PATH", "")
os.add_dll_directory(r"C:\My_Project\src\deeplearning\SAM3\TensorRT-RTX-1.5.0.114\bin")
os.add_dll_directory(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin")
# 使用标准 TRT 时还需注册插件
import tensorrt as trt
trt.init_libnvinfer_plugins(trt.Logger(trt.Logger.WARNING), "")
注意 :TRT-RTX 导入 tensorrt_rtx 模块,标准 TRT 导入 tensorrt 模块。
4. ONNX Runtime
用于预处理(图像缩放、归一化)。项目使用 onnxruntime-win-x64-gpu-1.26.0。
C:\My_Project\src\deeplearning\SAM3\sjl\onnxruntime-win-x64-gpu-1.26.0\onnxruntime-win-x64-gpu-1.26.0\
PyPI 安装:
bash
pip install onnxruntime-gpu==1.26.0
PyTorch 自带 cuDNN(版本 91002),独立于系统 cuDNN,不影响 TensorRT。
5. Python 环境
项目使用 Anaconda 管理环境:
bash
conda create -n sam3 python=3.10
conda activate sam3
pip install numpy opencv-python onnxruntime-gpu tensorrt-cu12
二、ONNX 模型导出
SAM3 ONNX 导出采用两步法:
Step 1: 导出核心模型
位置:sjl/export_onnx/export_onnx/export_onnx.py
由于 F.interpolate(antialias=True) 无法直接导出为 ONNX,预处理被剥离出核心模型。核心模型只包含 ViT 编码器 + 解码器 + 分割头。
Step 2: 导出预处理模型
位置:sjl/export_onnx/export_onnx/export_preprocess_onnx.py
用 ONNX 原生算子构建预处理流程:BGR→RGB 转换 → Resize(1008×1008) → 归一化。
导出产物在 sjl/export_onnx/exports/ 目录下,包含 FP16 和 FP32 两个版本。
多类别文本检测
SAM3 原生支持批量文本编码(forward_text() 接受 List[str])。
- 方案 A(推荐):导出时预先编码所有类别,作为 ONNX 常量。推理时无文本编码开销,但类别在导出时固定。
- 方案 B :将
language_features和language_mask设为 ONNX 输入,运行时灵活但需 CPU 预计算文本特征。
三、TensorRT Engine 构建与测试
Engine 文件列表
| 文件 | 来源 | 精度 | 大小 |
|---|---|---|---|
sam3_core_fp32_trt_rtx.engine |
TRT-RTX | FP32 | ~1.9G |
sam3_core_fp16_trt_rtx.engine |
TRT-RTX | FP16 | ~1.0G |
sam3_core_bf16_trt_rtx.engine |
TRT-RTX | FLOAT (非 BF16) | ~1.9G |
sam3_core_fp16_from_fp32.engine |
TRT 10.11, --fp16 |
FP16 | ~1.0G |
sam3_core_fp16_notf32.engine |
TRT 10.11, --fp16, --noTF32 |
FP16 | ~1.0G |
测试脚本
| 脚本 | 用途 |
|---|---|
test_trt_rtx.py |
主测试脚本(TRT-RTX 推理 + 计时 benchmark) |
test_fp16_quick.py |
快速 FP16 验证 |
profile_fp16.py |
FP16 vs FP32 精度对比 |
dump_layers.py |
层结构分析 |
dump_layers2.py |
层名称逐行 dump |
推理流程
原始图像 (cv2.imread)
→ ONNX 预处理 (CPU) → (1, 3, 1008, 1008) FP32/FP16
→ CUDA memcpy → GPU
→ TensorRT engine 推理 → 输出 (boxes, scores, masks, presence_logit)
→ CUDA memcpy → CPU
→ 后处理 (NMS, 阈值)
四、FP16 精度问题排查结果
问题现象
FP16 engine 推理得到的 detection scores 全部异常偏低(max=0.02),而期望值应接近 0.94。
| Engine | presence_logit | scores max | 精度正确? |
|---|---|---|---|
| ONNX FP32 CPU | +5.67 | 0.942 | 正确 |
| ONNX FP16 CPU | +5.67 | 0.942 | 正确 |
| TRT-RTX FP32 | +5.67 | 0.942 | 正确 |
| TRT-RTX FP16 | -2.82 | 0.021 | 错误 |
| TRT 10.11 FP16 | -4.66 | 0.003 | 错误 |
排查步骤
| # | 测试 | 结果 | 结论 |
|---|---|---|---|
| 1 | ONNX FP16 CPU 推理 | 正确 | FP16 模型本身没问题 |
| 2 | TRT-RTX FP32 GPU 推理 | 正确 | GPU 执行路径正常 |
| 3 | TRT-RTX FP16 GPU 推理 | 错误 | FP16 GPU kernel 有问题 |
| 4 | 升级 cuDNN 8.9.7 → 9.23.0 | 无变化 | 不是 cuDNN 单版本问题 |
| 5 | TRT 10.11 FP16 | 错误 | 不是 TRT-RTX 特有的 bug |
| 6 | --builderOptimizationLevel=0 |
错误 | 不是优化级别导致 |
| 7 | --noCompilationCache |
错误 | 不是编译缓存导致 |
| 8 | --noTF32 |
错误 | 不是 TF32 精度模式导致 |
速度对比(N=30 次迭代)
| Engine | 平均/帧 | FPS |
|---|---|---|
| TRT-RTX FP32 | 509ms | 2.0 |
| TRT-RTX FP16 | 122ms | 8.1 |
| TRT 10.11 FP16 | 95ms | 10.4 |
FP16 速度快 4-5 倍,但精度完全不可用。
层结构分析
- FP32 engine:1042 层
- FP16 engine:786 层(FP16 模式有更多层融合)
- 基本层类型相同,仅融合程度不同
- 完整层 dump 见
trt_test/layers_dump.txt
初步怀疑-待验证
TensorRT 的 FP16 GEMM/卷积 GPU kernel 在 Blackwell SM_120 架构上计算错误。
关键证据:
- ONNX FP16 CPU 正确 → 模型 FP16 精度本身够用
- TRT FP32 GPU 正确 → GPU 执行正常
- TRT-RTX 和 TRT 10.11 都错 → 不是某个 TRT 版本的 bug
- cuDNN 升级无效 → 不是单一 cuDNN 版本问题
SM_120 是 2025 年新架构,TensorRT/cuDNN 的 FP16 kernel 在该架构上未经验证。其他项目(如 dataplayer12/SAM3-TensorRT)使用 RTX 3090 (Ampere)、A100、H100、GB200 (服务器 Blackwell) 等成熟 GPU,FP16 正常工作,问题仅出现在桌面版 SM_120a (RTX 5060 Ti)。
五、解决方案与下一步
当前可用方案
使用 FP32 engine:
- 精度:正确(scores max=0.942,与 ONNX 一致)
- 速度:509ms/帧,2 FPS
- Engine 文件:
sam3_core_fp32_trt_rtx.engine
修复方向
- **使用更新版本的 cuDNN **:后续版本可能修复 SM_120 FP16 kernel
- 等待 TRT-RTX 2.x 或 TRT 11.x:新 TRT 版本可能加入 SM_120a 验证
- CUDA 12.9 完整工具链:NVIDIA 发布正式版 CUDA 12.9 后,可能包含 SM_120 修复
- 测试 BF16 路径 :如果 NVIDIA 在 TRT-RTX 中加入
--bf16支持,BF16 可能走不同的 GPU kernel 路径