初探SAM3(一) :SAM3 TensorRT 部署踩坑记录

本文档记录在 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_featureslanguage_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 架构上计算错误。

关键证据:

  1. ONNX FP16 CPU 正确 → 模型 FP16 精度本身够用
  2. TRT FP32 GPU 正确 → GPU 执行正常
  3. TRT-RTX 和 TRT 10.11 都错 → 不是某个 TRT 版本的 bug
  4. 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

修复方向

  1. **使用更新版本的 cuDNN **:后续版本可能修复 SM_120 FP16 kernel
  2. 等待 TRT-RTX 2.x 或 TRT 11.x:新 TRT 版本可能加入 SM_120a 验证
  3. CUDA 12.9 完整工具链:NVIDIA 发布正式版 CUDA 12.9 后,可能包含 SM_120 修复
  4. 测试 BF16 路径 :如果 NVIDIA 在 TRT-RTX 中加入 --bf16 支持,BF16 可能走不同的 GPU kernel 路径
相关推荐
marsh02061 小时前
61 openclaw电商系统架构:从需求到实现的完整方案
ai·系统架构·技术美术
元拓数智1 小时前
NL2SQL落地后的“隐形陷阱”:如何破解数据变更引发的查询失效难题
大数据·ai·自然语言处理·nl2sql
云草桑2 小时前
.NET10+AI 架构师全套实战学习文档(含源码、案例、面试题、项目源码)
人工智能·学习·ai·.net
器灵科技2 小时前
DeepSeek V4 Pro宣称:超GPT-5.5+永久降价75%
大数据·人工智能·gpt·阿里云·ai·语言模型
yyoc972 小时前
本地 Flink on K8s + Iceberg + MinIO 实时数仓平台 — AI部署指南与踩坑实录
大数据·ai·flink·kubernetes·iceberg
蜂蜜黄油呀土豆2 小时前
ReWOO 与 Plan-and-Execute:解耦的规划
python·ai·大模型
金融RPA机器人丨实在智能2 小时前
工程线索工具合规避坑指南:使用开源爬虫抓取数据会触犯法规吗?实在Agent给出了安全答案
人工智能·爬虫·安全·ai·开源
Java知识技术分享3 小时前
node安装新版本,并解决opencode和claude code不能用问题
ai·个人开发·ai编程
我认不到你3 小时前
【开源、教程】RAG全流程实现(java+完整代码):第一弹
java·开发语言·人工智能·深度学习·ai·语言模型·开源