初探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 路径
相关推荐
风fffff3 小时前
dsh-project-memory v0.3.0到v0.4.0:从项目记忆到开发工作流记忆的演进
ai·agent·插件·dsh·deepseek harness
敢敢是只喵i3 小时前
给现有业务系统接 AI 助手,入口怎么选?网页聊天、自动化任务与本地 Agent 对比
运维·人工智能·ai·开源·自动化·安全架构
ddshub_cc4 小时前
GPT Image 2 电商场景指南:主图、详情页与活动物料怎么出
gpt·ai·image2·ai生图·电商ai·电商工具·电商生图
MinggeQingchun4 小时前
AI - 阿里云百炼
ai·阿里云百炼
章老师说4 小时前
BFE v1.8.6 正式发布:AI 网关计费精细化、Claude 协议与会话亲和性升级
运维·人工智能·ai·负载均衡·ai-native
ai小陈5 小时前
FramePack图生视频云端部署实战:从单图输入到视频输出的完整流程
服务器·人工智能·安全·ai·音视频·gpu算力
长沙京卓5 小时前
Copilot Coding Agent 变了:AI 编程正在从插件变成项目成员
人工智能·ai
is me monday5 小时前
Codex-CLI 下载安装与 Token 配置教程(附codex安装包)
ai
ai小陈6 小时前
Hunyuan3D-2云端部署实战:图生3D、文生3D怎么跑更稳
人工智能·科技·3d·ai·音视频·gpu算力
村口徐大爷7 小时前
电商新媒体通用!Lingko AI全套视觉素材自动化量产落地实操教程
人工智能·ai·ai工具·电商运营·lingnko ai