简介面向C与AI推理开发者的SAM模型TensorRT部署方案提供完整C源码与部署步骤解决Segment Anything在工业视觉场景中的高性能落地问题。资源共22个文件压缩包仅1.74MB以C头文件与源文件为主辅以Jupyter Notebook导出教程、JSON/CMake配置、Markdown使用文档以及示例图片与动图覆盖模型转换、推理加速、接口封装等关键环节。已有815人学习下载适合具备一定CUDA/TensorRT基础的CV工程师快速上手也适用于图像分割、目标检测等需要高效推理的视觉任务。内容包含SAM模型导出与优化脚本、C推理主程序、线程池及公共工具模块并附Windows部署文档与README读者可依据目录结构从环境配置到运行演示逐步搭建有效缩短工程落地周期。1. 用 TensorRT 把 SAM 从“演示玩具”变成可交付的 C 服务SAMSegment Anything“分割一切”是个通过点、框、掩码提示做通用分割的大模型但它在 PyTorch 里跑想法和在企业内网 GPU 服务器上做本地部署是两码事前者能忍受几秒推理后者要面对 C 服务、并发请求、显存和稳定延迟。TensorRT 是英伟达把这类 ViT 编码器压得最狠的推理引擎配合 C 源码把图像编码器、提示编码器、掩码解码器串成一个可执行程序才能让“分割一切”真正接进业务。这篇文章按模块拆分、ONNX 导出、TensorRT 转引擎、C 加载运行这条链路讲并把真实部署里踩过的坑摊开。适合准备把 SAM 集成进 C 项目的工程师也适合想评估这条路径值不值得投入的技术负责人。2. 拆解 SAM 的三个子模块把算力花在明白处再谈 TensorRT 加速拿到一份 SAM 的 C 部署源码我第一件事不是开 IDE而是先看模型怎么分装。SAM 对外是“一张图进去、一个掩码出来”实际内部是三个耦合很松的模块图像编码器、提示编码器、掩码解码器。搞清楚每个模块的输入输出和精度敏感度直接决定了你后面 ONNX 导出怎么切图、trtexec 怎么配参数。很多新手上来就把整个 SAM 导成一个 ONNX结果又大又慢就是这个环节没想明白。2.1 图像编码器占掉八成以上算力固定尺寸下的 ViTSAM 的图像编码器是一个视觉 TransformerViT变体把 1024×1024 的输入图像编码成一张 256 通道的图像嵌入image embedding。这个 1024×1024 是 SAM 官方的标准输入尺寸不是链路里随便选的换尺寸意味着位置编码和 patch 切分全部重算模型行为也会跟着变。这一段的计算量在整个 SAM 推理里占绝对大头权重也有几百 MB 的量级。提示编码器和掩码解码器加起来都不到它的零头所以部署时算力必须优先砸在图像编码器上。TensorRT 对固定形状的 engine 能做的优化非常激进把卷积、归一化、注意力里的算子做融合提前选好 kernel显存 buffer 复用。反过来一旦输入尺寸是动态的很多优化就锁死了。所以我的第一个选择是图像编码器单独导出一个固定形状的 engine输入严格锁成 1×3×1024×1024。很多源码包也默认这么拆你在src/里看到两个 engine 的加载逻辑而不是一个就是这个原因。2.2 提示编码器与掩码解码器动态点数与轻量化模块提示编码器把用户给的稀疏提示和稠密提示变成 embedding。点坐标要跟位置编码叠加框要编码成左上右下两个点掩码输入要先过一个小卷积网络。这个模块本身很小但它有个部署最烦的属性点数是动态的。用户一次点一个点还是五个点ONNX 模型里对应的坐标张量维度就不同。掩码解码器是个轻量 Transformer输入是图像 embedding 加上 prompt embedding输出低分辨率掩码和置信度分数再做上采样回到原图尺寸。模型尺寸小但数值敏感度高因为它直接决定掩码边界。这也是为什么我会把“提示编码器掩码解码器”拆成第二个 engine并特意把 num_points 这个维度做成动态图像编码器固定形状求极致性能提示解码器用动态形状换取灵活性。两者合一当然也能跑但动态轴会污染图像编码器的优化空间而且每次换提示都要连带重算图像编码器工程上很不划算。2.3 为什么我不直接用 ONNX Runtime 或 OpenVINO 顶替 TensorRT经常有人问SAM 不也能用 ONNX Runtime 跑吗为什么非要 TensorRT我的回答是看交付环境。ONNX Runtime 的 CUDA EP 对常见 CNN 表现不差但对 ViT 里的多头注意力、GELU、层归一化这些算子的融合收益明显不如 TensorRT它的强项是跨硬件和快速迭代适合原型验证和低并发。OpenVINO 在 Intel 平台上挺能打但你在 N 卡上跑等于绕了远路。TensorRT 的路子是事前做一遍 kernel 选择和图优化生成针对你这块 GPU 的 engine跑起来之后算子边界少、显存复用更彻底。对企业大模型私有化部署、统一用 N 卡的内网环境来说这是最顺的一条链路代价是版本矩阵敏感、调试黑匣子比较深后面第 4 章的坑基本都是这么来的。如果你只是临时验证效果ONNX Runtime 当然更省事但一旦进入服务化发布、响应时间有要求的阶段我会把宝压在 TensorRT 上。3. 从 zip 到可运行 C 程序源码编译与最小部署链路源码包解压以后我习惯先扫一眼目录里有没有 README、CMakeLists.txt 和 scripts/。如果你的包是我常见的那种结构——src/ 放 C 源码scripts/ 放导出和转换脚本third_party/ 放依赖——那大致的链路就是PyTorch 权重导出 ONNX → trtexec 转 engine → cmake 构建 C 工程 → 跑最小 demo。下面按这条链路逐步过每一步我给参数和选择逻辑避免你照着敲完发现死在版本上。3.1 环境匹配先于编译CUDA、cuDNN、TensorRT 的版本对齐第一个容易翻车的地方不是代码是环境。TensorRT 的 C API 链接的是 libnvinfer.so它和 CUDA、cuDNN 的版本非常敏感。我先跑几条命令确认底数nvidia-smi # 看驱动和 GPU 型号 nvcc --version # 看 CUDA 工具链版本 ls /opt/TensorRT*/lib # 看 TensorRT 库文件存在版本匹配我一般这样对齐| 组件 | 推荐基线 | 说明 | | CUDA | 12.212.4 | 以驱动支持为准不追最新 | | cuDNN | 8.9.x | TensorRT 部分算子依赖 cuDNN | | TensorRT | 8.6 / 10.x 稳定版 | 避开刚发布的 .0 版本 | | 编译环境 | GCC 9 / MSVC 2019 | Windows 上记得装 Visual C Redistributable |如果你链接完跑起来报undefined symbol先别怀疑代码用ldd看可执行文件到底链接了哪个 libnvinferldd bin/sam_demo | grep nvinfer再把两边的版本号对一眼八成是 CUDA 或 cuDNN 的小版本不对齐。这类问题花一晚上查源码不如先花十分钟对版本。3.2 从 PyTorch 权重导出 ONNX图像编码器固定形状解码器放开动态轴假设你已经有一个 SAM 的 PyTorch checkpoint。导出核心思路是分块导出图像编码器固定形状提示编码器加掩码解码器放开动态点。下面是图像编码器的导出写法import torch from segment_anything import sam_model_registry checkpoint sam_vit_h_4b8939.pth model sam_model_registry[vit_h](checkpointcheckpoint).cuda().eval() image torch.randn(1, 3, 1024, 1024).cuda() torch.onnx.export( model.image_encoder, image, sam_image_encoder.onnx, input_names[image], output_names[image_embedding], opset_version17, do_constant_foldingTrue, )opset_version17是我常用的起点低于 14 有些动态 shape 表达式可能导不出来过高版本则可能引入 TensorRT 解析器不认识的算子。do_constant_foldingTrue可以把只跟权重相关的计算提前折叠缩减 ONNX 节点数对后面转 engine 有好处。提示解码器这一侧点坐标必须动态。如果写死成单个点用户点两个点就崩。常见做法是把两个模块包成一个 wrapper只放开 num_points 这个轴class DecoderWrapper(torch.nn.Module): def __init__(self, model): super().__init__() self.prompt_encoder model.prompt_encoder self.mask_decoder model.mask_decoder def forward(self, image_embedding, point_coords, point_labels): sparse, dense self.prompt_encoder( points(point_coords, point_labels), boxesNone, masksNone, ) masks, scores self.mask_decoder( image_embeddingsimage_embedding, image_peself.prompt_encoder.get_dense_pe(), sparse_prompt_embeddingssparse, dense_prompt_embeddingsdense, multimask_outputTrue, ) return masks, scores torch.onnx.export( decoder_wrapper, (image_embedding, point_coords, point_labels), sam_prompt_decoder.onnx, input_names[image_embedding, point_coords, point_labels], output_names[masks, scores], dynamic_axes{ point_coords: {1: num_points}, point_labels: {1: num_points}, }, opset_version17, )这里把multimask_outputTrue写死一次性输出三个候选掩码和对应分数C 端直接挑分最高的。动态轴只放开 num_points图像 embedding 保持固定不用把 batch 也放开SAM 的部署大多一次处理一张图没必要给自己找麻烦。负样本点也是用 point_labels 里的 0 和 1 表达的导出时不需要额外结构。3.3 用 trtexec 生成 Engine三条命令和常用参数取舍ONNX 只是中间产物真正给 C 用的是 TensorRT engine 文件。我直接用 trtexec 转不额外写 Python 脚本trtexec --onnxsam_image_encoder.onnx \ --saveEnginesam_image_encoder.engine \ --fp16 \ --minShapesimage:1x3x1024x1024 \ --optShapesimage:1x3x1024x1024 \ --maxShapesimage:1x3x1024x1024 trtexec --onnxsam_prompt_decoder.onnx \ --saveEnginesam_prompt_decoder.engine \ --fp16 \ --minShapesimage_embedding:1x256x64x64:point_coords:1x1x2:point_labels:1x1 \ --optShapesimage_embedding:1x256x64x64:point_coords:1x5x2:point_labels:1x5 \ --maxShapesimage_embedding:1x256x64x64:point_coords:1x16x2:point_labels:1x16图像编码器这条三组 shape 完全相同等于告诉 TensorRT 按固定 shape 优化kernel 选择空间更大。解码器这条把优化点定在 5 个点最大放开到 16 个点如果业务里用户最多点 10 个就把 maxShapes 写成 10因为 max 越大TensorRT 预留的显存和 kernel 分支也越多。--fp16开不开取决于精度要求。图像编码器开 FP16 基本无损但解码器最好先做对比再定后面 4.3 会细说。trtexec 跑完会输出一行Average latency那个数字只能当参考真实服务延迟要看 C 端在并发下的表现。3.4 C 侧最小加载流程从 engine 落地到一次完整推理生成完 engineC 端要做的事很固定读引擎文件 → 反序列化出 ICudaEngine → 创建 IExecutionContext → 绑定输入输出缓冲 → enqueue。下面是去掉业务细节的最小骨架#include NvInfer.h #include fstream #include vector std::vectorchar loadFile(const std::string path) { std::ifstream f(path, std::ios::binary); return {std::istreambuf_iteratorchar(f), std::istreambuf_iteratorchar()}; } int main() { auto runtime std::unique_ptrnvinfer1::IRuntime( nvinfer1::createInferRuntime(gLogger)); auto engineData loadFile(sam_image_encoder.engine); auto engine std::unique_ptrnvinfer1::ICudaEngine( runtime-deserializeCudaEngine(engineData.data(), engineData.size())); auto context std::unique_ptrnvinfer1::IExecutionContext( engine-createExecutionContext()); void* inputDevice nullptr; void* outputDevice nullptr; cudaMalloc(inputDevice, 1 * 3 * 1024 * 1024 * sizeof(float)); cudaMalloc(outputDevice, 1 * 256 * 64 * 64 * sizeof(float)); context-setTensorAddress(image, inputDevice); context-setTensorAddress(image_embedding, outputDevice); cudaStream_t stream; cudaStreamCreate(stream); context-enqueueV3(stream); cudaStreamSynchronize(stream); return 0; }注意这里用的是 TensorRT 8.5 之后的setTensorAddress enqueueV3新接口比老的enqueueV2少了一堆 binding index 换算代码干净很多。如果你是旧版 TensorRT就得退回enqueueV2参数也改成setBindingAddress那套写法。这段代码只够验证链路通不通。工程上常见做法是把它包成一个SamEngine类图像编码器只跑一次image embedding 存下来给多个 prompt 请求复用提示解码器每次请求用独立 context。图像编码器的输出是 1×256×64×64 的四维张量这个尺寸在三段链路里基本是硬编码的换 backbone 时记得同步改。4. 部署避坑我在 SAM 的 TensorRT 实践中跳过的五个坑这一章都是实际踩过、或者看别人踩完来求救的案例每条按现象 → 原因 → 解决写查问题时可以直接按图索骥。4.1 掩码错位甚至全空预处理与训练时不一致现象点一个位置返回的掩码要么全黑要么明显偏到别的地方。原因SAM 官方预处理不是简单cv2.resize成 1024×1024而是先把长边缩放到 1024再对短边做填充最后按 ImageNet 的 mean/std 归一化。很多人图省事直接拉伸图片长宽比一变ViT 的位置编码空间就歪了掩码当然对不上。还有一版错误是先把像素除以 255再用原始 mean/std 做归一化等于两套尺度混用。解决在 C 端按 SAM 的预处理逻辑重写等比缩放、填充、归一化一次到位。归一化这一步可以直接对 0255 的像素做(pixel - mean) / stdfloat mean[3] {123.675f, 116.28f, 103.53f}; float std[3] {58.395f, 57.12f, 57.375f}; // 输入已resize并pad到1024x1024RGB顺序 // out[c][h][w] (pixel[c][h][w] - mean[c]) / std[c]常用的校验手段是用同一张图分别跑 PyTorch 原模型和 C 端比较 image embedding 的余弦相似度低于 0.99 就是预处理不一致。4.2 一进动态形状推理就报错context 的输入维度没有显式设置现象engine 构建成功跑固定点数也正常一旦换不同点数推理就报Assertion failed或者返回空数据。原因解码器的 ONNX 里 point_coords 是动态轴但 C 端没有在每次推理前告诉 context 这次的维度。TensorRT 的 context 默认沿用上一次的形状第一次是 1 个点第二次来了 5 个点它并不知道。解决每次 enqueue 之前显式设置输入形状nvinfer1::Dims4 dims; dims.nbDims 3; dims.d[0] 1; // batch固定 dims.d[1] numPoints; // 本次请求的点数 dims.d[2] 2; // x, y 坐标 context-setInputShape(point_coords, dims);记住profile 里配的 maxShapes 是上限不代表 context 会自动跟随输入变化。如果用的是旧接口还要同时把输出 binding 的维度重新setBindingDimensions。4.3 FP16 让你的掩码肉眼可见地变烂解码器别跟风压精度现象图像编码器开 FP16 没问题但整条链路 mask 质量明显下降小目标掩码容易破边界发毛。原因分割任务的输出层对数值精度比较敏感掩码解码器里的 sigmoid、上采样和注意力叠加FP16 会把小数值的区分度压没。解决分开转换、分开压精度。图像编码器用--fp16提示解码器保持 FP32。两个 engine 独立加载解码器这边推理一次也就几毫秒损失一点速度换掩码质量非常值得。如果要上 INT8SAM 这种大模型需要逐层校准校准集选不好掩码基本废了我不建议新手直接碰。4.4 服务跑一段时间段错误context 被并发复用现象demo 单线程正常一上多线程并发就随机段错误或者多个请求交错后输出错乱。原因同一个 IExecutionContext 被多个线程同时拿来 enqueue。TensorRT 的 context 内部有状态输入 shape、binding 地址都是上下文相关的并发复用等于把黑匣子当普通函数调。解决一个 context 只服务一个线程。engine 是只读的可以全局共享context 按线程或按请求创建做完回收。以 8 并发为例创建 8 个 context 放进对象池任务进来取一个用完归还。context 创建开销远小于重建 engine这种设计既控制显存又好排查。4.5 权重文件缺失或加载报 key 错误先把模型文件对齐再谈转换现象源码包拿到手sam_model_registry[vit_h](checkpoint...)直接报 key 不匹配或者文件下载到一半失败。原因SAM 有 vit_b、vit_l、vit_h 几个 backbonecheckpoint 和注册的模型类型必须严格一致。打包的人经常把权重文件省略掉只给脚本默认你能从外网下下来。解决先把 checkpoint 放到位并且确认文件名、模型类型一一对应。公司内网不能下载时在 Python 环境跑一次加载确认 key 对得上再导出。我见过有人拿 vit_b 的权重加载 vit_h 的模型转换全程无报错结果跑出来全是噪点查了半天才发现是权重版本错了。5. 从能跑到能发布验证确定性、并发 context 和 CUDA Graph 优化最后聊三个我每次部署 SAM 都会做的动作不复杂但能把“demo 能跑”变成“服务能上线”。顺序不要反先验证再并发最后优化。验证。用 PyTorch 原模型输出当基准上线前做一次持久的盯梢检查取一张图分别跑两端比较 image embedding 的余弦相似度我要求稳定在 0.99 以上。低于这个数是预处理或权重有问题别急着调并发。并发。开业之前把 context 放进对象池。engine 反序列化一次进程全局共享context 按需创建用完归还。注意给池子加锁线程取用还是放回都要抢一把互斥锁class SamEnginePool { std::mutex mu_; std::dequestd::unique_ptrnvinfer1::IExecutionContext pool_; public: std::unique_ptrnvinfer1::IExecutionContext acquire() { std::lock_guardstd::mutex lock(mu_); if (pool_.empty()) { pool_.push_back(engine_-createExecutionContext()); } auto ctx std::move(pool_.front()); pool_.pop_front(); return ctx; } void release(std::unique_ptrnvinfer1::IExecutionContext ctx) { std::lock_guardstd::mutex lock(mu_); pool_.push_back(std::move(ctx)); } };优化。TensorRT 的enqueueV3很适合配合 CUDA Graph 使用把解码器的推理录制成一次图重放开销很低。我的做法是对提示解码器这种高频小模型录图图像编码器这种大模型反而收益不明显因为瓶颈在访存和带宽。记录一次之后反复cudaGraphLaunch单次推理的启动开销能压下去不少。我以前的教训是图省显存让多个请求共享一个 context结果 CUDA Graph 重放时把 stream 搞乱掩码错得一塌糊涂查了一下午才发现是 context 被并发复用。现在我的习惯是engine 尽量共享context 一个请求一个显存不会爆问题也好定位。希望帮到你。本文还有配套的精品资源点击获取