简介：面向需要将 Segment Anything Model 落地到 NVIDIA GPU 的算法工程师与 C++ 开发人员，这套资源完整给出 TensorRT 部署 SAM 分割模型的工程代码与分步部署流程。内容覆盖模型转换、层融合、内核自动调优、推理执行等关键环节，适合已有 PyTorch 基础、希望掌握 TensorRT 推理优化或正在搭建分割服务的读者参考。压缩包共 29 个文件、约 5.32MB，以 .h/.cpp 源码、CMakeLists 构建配置、.ipynb 教程与 .md 说明文档为主，另含 Dockerfile.dev、ThreadPool、VSCode 配置及可视化样例，便于对照源码理解 API 调用与执行链路，也方便在容器或本地环境快速复现。资源中附带了 README 和 Windows 部署说明，可帮助使用者完成模型导出、优化、推理验证等环节；内含的教程文档和效果演示能进一步降低门槛，支撑后续迁移到实际项目。该资源已有 245 人学习，虽体量不大，但代码与文档组织清晰，是一份聚焦 TensorRT 推理落地的实用参考。
1. TensorRT部署SAM分割模型：为什么C++落地比Python快一个量级
做图像分割的同行应该都有这种感觉：SAM（Segment Anything Model）在Python里跑Demo很爽，但要上生产、接实时视频流、塞进边缘盒子，立刻就被PyTorch的显存占用和推理延迟卡住。TensorRT部署SAM分割模型的C++代码及部署流程，本质上就是把SAM的Image Encoder、Prompt Encoder、Mask Decoder三段拆开，分别转成TensorRT引擎，再用C++写一个推理管线把它们串起来。这个方案能解决的是：模型还在，精度不掉太多，但显存占用能降一半以上，单帧推理延迟从几百毫秒压到几十毫秒，而且彻底摆脱Python运行时依赖。适合谁？手里有SAM权重、想把它做成C++ SDK或者嵌入现有服务的人。我建议你先别急着抄代码，先想清楚一件事：SAM不是单一输入输出的黑匣子，它的输入有图像和Prompt两种，输出是掩码和分数，部署时要拆解成多个引擎，否则后面每一步都在填坑。
2. 先拆清SAM三段架构：TensorRT化之前必须懂的边界
2.1 Image Encoder是重头戏：ViT/BEVT的算子映射与固定输入尺寸
SAM的Image Encoder是一个基于ViT的骨干网络，输入是1024×1024的RGB图像，输出是64×64×256的嵌入向量。TensorRT化最大的坑在于ViT里的多层多头注意力、LayerNorm、SiLU这些算子在TRT不同版本上的支持程度不同。常见做法是先把PyTorch模型导成ONNX，再转TRT。但SAM官方给的原始权重是ViT-H/ ViT-B，导ONNX时动态轴问题就来了——SAM的Image Encoder内部是固定分辨率处理的，但外围的Resize逻辑在Python里做的。我一般会直接固定输入尺寸为1×3×1024×1024，避免动态shape带来的额外优化损失。参数上注意：TRT的builder里设置setMaxBatchSize(1)就够了，batch=1的实时场景没人会用batch=8跑分割。还有setFp16Mode(true)，SAM的Encoder在FP16下精度损失很小，Mask边界会有轻微抖动，但可控。如果你是做医学影像这类对边界敏感的场景，建议保持FP32，或者用INT8+校准集，否则后面边界像素的误判会让你怀疑人生。
2.2 Prompt Encoder与Mask Decoder：别一把梭哈成一个引擎
提示编码器接收点、框、掩码三类提示，输出稀疏和稠密嵌入；掩码解码器把嵌入和图像特征融在一起，输出三个分辨率的掩码和IoU预测。这两个模块网络浅、计算量小，但输入结构五花八门——点坐标是动态的、框的数量不固定，如果强行和Encoder合并成一个TRT引擎，你会被动态shape的优化空间搞到崩溃。实际部署时，我会把它们拆成两个小引擎，或者干脆用ONNX Runtime跑这两个子模块，只在Encoder上用TRT，因为Encoder占了SAM约95%的计算量。这个取舍很重要：TensorRT部署SAM分割模型C++代码，核心收益几乎全部来自Encoder加速，Decoder用TRT收益很小，却会引入很多动态shape的麻烦。
2.3 三段拼接的TensorRT集成方案对比
| 方案 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| 整体导出为ONNX→TRT | 管线简单 | 动态Prompt处理差，优化难 | 固定单框、离线批量 |
| Encoder走TRT，Decoder用ONNX Runtime | 灵活，开发快 | 两套runtime，内存管理繁琐 | 实时交互式分割 |
| Encoder/Decoder都走TRT，用C++封装 | 延迟最低，显存可控 | 需要处理动态shape和显存复用 | 生产级SDK、边缘部署 |
我自己推荐第三种，但工程量大很多。前两种适合原型验证。真正投产时，第三种才值得花功夫，因为TensorRT的精髓就在引擎内部做了层融合和显存池化，只加速一半等于没吃到全部红利。下面开始讲我验证过的具体路径。
3. 从ONNX导出到TensorRT引擎：最小可跑的C++推理管线
3.1 导出固定shape的ONNX，避开动态轴这个坑
先用Python把SAM的Image Encoder单独导成ONNX。关键代码：
import torch from segment_anything import sam_model_registry sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth").cuda().eval() # 只取ImageEncoder，不要整个SAM image_encoder = sam.image_encoder dummy = torch.randn(1, 3, 1024, 1024).cuda() with torch.no_grad(): torch.onnx.export( image_encoder, dummy, "sam_encoder.onnx", opset_version=17, input_names=["image"], output_names=["image_embedding"], dynamic_axes=None, # 固定shape，避免动态轴 )这里必须设置dynamic_axes=None，否则导出会生成动态shape的ONNX，后面TRT转换时虽然能转，但会禁用很多优化，速度反而不如静态。opset_version建议用17或18，高版本对LayerNorm、Gelu的支持更完整。别忘了torch.onnx.export默认会走trace，SAM里的插值操作和mask_decode无关，但Encoder内部有window_position这种张量操作，固定输入尺寸100%避免trace出错。
导出后用onnxsim精简一下模型，能干掉一部分冗余Reshape和Transpose。
3.2 用trtexec生成FP16引擎，先跑通再写代码
TRT引擎生成最快的方式是用trtexec命令行，不用先写C++。命令如下：
/opt/TensorRT/bin/trtexec \ --onnx=sam_encoder.onnx \ --saveEngine=sam_encoder_fp16.engine \ --fp16 \ --minShapes=image:1x3x1024x1024 \ --optShapes=image:1x3x1024x1024 \ --maxShapes=image:1x3x1024x1024 \ --workspace=4096参数说明：--fp16开启半精度，--workspace给TensorRT分配4GB显存做 tactic 搜索，实际运行时workspace不占用这么多，但选tactic时需要。--minShapes/optShapes/maxShapes虽然这里是静态shape，但写上能防止某些算子故意生成动态版本。跑通后trtexec会打印Throughput和Latency，先记录FP32和FP16的对比。我见过很多人在这一步省事不跑trtexec，直接写C++调deserializeCudaEngine，结果引擎加载失败都不知道是CUDA版本不匹配还是ONNX里有TRT不支持的算子。trtexec是唯一的排错镜。
3.3 C++推理类的最小实现：加载引擎与执行上下文
引擎生成后，写一个C++类封装加载、推理、后处理。最小骨架如下：
#include <NvInfer.h> #include <cuda_runtime_api.h> #include <fstream> #include <vector> class SamEncoder { public: bool loadEngine(const std::string& path) { std::ifstream file(path, std::ios::binary); file.seekg(0, std::ios::end); size_t size = file.tellg(); file.seekg(0, std::ios::beg); std::vector<char> blob(size); file.read(blob.data(), size); // 创建runtime和engine runtime_ = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger_)); engine_ = std::unique_ptr<nvinfer1::ICudaEngine>(runtime_->deserializeCudaEngine(blob.data(), size)); context_ = std::unique_ptr<nvinfer1::IExecutionContext>(engine_->createExecutionContext()); return context_ != nullptr; } bool infer(float* input, float* output) { void* buffers[2]; cudaMalloc(&buffers[0], 1 * 3 * 1024 * 1024 * sizeof(float)); cudaMalloc(&buffers[1], 1 * 256 * 64 * 64 * sizeof(float)); cudaMemcpy(buffers[0], input, 1 * 3 * 1024 * 1024 * sizeof(float), cudaMemcpyHostToDevice); bool ok = context_->enqueueV2(buffers, 0, nullptr); cudaMemcpy(output, buffers[1], 1 * 256 * 64 * 64 * sizeof(float), cudaMemcpyDeviceToHost); cudaFree(buffers[0]); cudaFree(buffers[1]); return ok; } private: std::unique_ptr<nvinfer1::IRuntime> runtime_; std::unique_ptr<nvinfer1::ICudaEngine> engine_; std::unique_ptr<nvinfer1::IExecutionContext> context_; nvinfer1::ILogger logger_; };这段代码逻辑是：先把.engine文件读进内存，反序列化成ICudaEngine，调用createExecutionContext生成执行上下文。推理时enqueueV2传入绑定的buffer指针数组。注意buffer顺序要和ONNX的输入输出顺序一致，这里输入是image，输出是image_embedding。生产环境不要每次推理malloc/free，应该用cudaMalloc一次，常驻显存。另外enqueueV2的stream参数传0是同步执行，实际部署要建一个cudaStream_t，用enqueueV3或者enqueueV2传入stream实现异步，否则CPU和GPU是串行的，延迟会虚高。
4. 部署流程的完整闭环：图像预处理、Prompt处理、掩码后处理
4.1 图像预处理：Resize、归一化，必须和Python版完全对齐
SAM在Python里的预处理是：先等比例缩放图像让长边=1024，再padding到1024×1024，最后做RGB均值/方差归一化。C++端如果用OpenCV实现，最容易翻车的是padding时填充值。PyTorch的transform里填充的是(0,0,0)，但归一化是mean=[123.675, 116.28, 103.53],std=[58.395, 57.12, 57.375]，像素值域是0-255。C++代码里要按这个顺序来：
cv::Mat resizeImg; float scale = 1024.0f / std::max(img.cols, img.rows); cv::resize(img, resizeImg, cv::Size(round(img.cols * scale), round(img.rows * scale))); cv::Mat padded(1024, 1024, CV_8UC3, cv::Scalar(0, 0, 0)); resizeImg.copyTo(padded(cv::Rect(0, 0, resizeImg.cols, resizeImg.rows))); // HWC -> CHW 并归一化 std::vector<float> input(1 * 3 * 1024 * 1024); #pragma omp parallel for for (int c = 0; c < 3; ++c) { for (int i = 0; i < 1024 * 1024; ++i) { input[c * 1024 * 1024 + i] = (padded.data[i * 3 + c] - mean[c]) / std[c]; } }这段代码里有个细节：padding区域参与了归一化，填充的0会被减均值除方差变成负值，这会影响Encoder的特征吗？实测影响很小，但如果你做的是细粒度分割，建议把padding区域也做成(114,114,114)的灰度填充，并记录pad的宽高，后处理时再裁掉。我在第一次部署时就是没对齐padding，导致生成的embedding和Python版对不上，Mask结果偏移了几个像素。
4.2 Prompt编码的C++实现与TRT引擎的输入输出约定
Prompt部分我拆成独立小引擎（或者直接CPU计算）。点坐标和框坐标在送到Encoder前，必须按原有预处理比例换算到1024×1024坐标系。比如原图1000×800，缩放后是1024×819.2，padding到1024×1024，点坐标要乘scale后加上左上偏移。这个换算错了，Prompt就点不到物体上。Mask Decoder的输入除了Prompt embedding外，还需要Image Encoder的embedding，所以C++管线里必须保持embedding的显存驻留，不能每次从GPU拷回CPU再传回去。我建议在GPU上分配embedding_buffer，Encoder输出直接留在显存，Decoder输入直接指向这个buffer，省一次H2D/D2H拷贝，这个优化能让端到端延迟减少2-3ms。
4.3 掩码后处理：从三个输出中选最优，并还原到原图尺寸
Mask Decoder输出三个不同尺度的掩码，每个掩码配上IoU预测分。C++里要做的：取IoU分数最高的那个掩码（或者用argmax策略），然后Sigmoid，阈值化到0.5，再按原图的缩放和padding还原。具体后处理代码：
// 假设mask是decoder输出，shape [1, 3, 256, 256]，scores是[1, 3] int bestIdx = std::max_element(scores, scores + 3) - scores; cv::Mat mask(256, 256, CV_32FC1, maskData + bestIdx * 256 * 256); cv::Mat maskBin; cv::threshold(mask, maskBin, 0.5, 1.0, cv::THRESH_BINARY); // 还原到1024x1024，再裁掉padding cv::resize(maskBin, maskBin, cv::Size(1024, 1024), cv::INTER_NEAREST); cv::Rect roi(0, 0, origWidth * scale, origHeight * scale); maskBin(roi).copyTo(finalMask);注意INTER_NEAREST，不要用线性插值，否则掩码边界会出现过渡像素。这里的scale和padding要和预处理严格对称，建议用同一个结构体记录，别在两处单独计算。
5. 避坑与常见问题：TensorRT部署SAM的5个血泪经验
5.1 引擎加载失败提示“Unsupported Layer”：多半是ONNX里混入了TRT不支持的算子
现象：deserializeCudaEngine或trtexec报错，常见于GridSample、Resize的坐标变换模式。
原因：SAM的Encoder里有F.grid_sample做窗口注意力中的位置偏移，ONNX导出后变成GridSample算子，低版本TensorRT不支持。解决：换用支持GridSample的TensorRT版本（8.6以上），或者把grid_sample替换成等价的affine_grid+bilinear sampling组合，但工作量大。我实际用的办法是用opset_version=17+TensorRT 8.6，顺利通过。如果你还在用TRT 8.2，建议先升级。
5.2 显存分配不足：输入分辨率固定了，Embedding buffer大小却算错
现象：推理时cudaMemcpy报invalid argument或者显存越界。
原因：Image Encoder输出是1×256×64×64的embedding，但如果你把输出Buffer分配成1×256×1024×1024，显存直接超。解决：打印引擎的binding维度，不要手算。代码里用engine_->getBindingDimensions(1)获取实际输出维度，再product计算总元素数。我犯过错，所以现在写死也要在加载引擎后校验一遍。
5.3 Prompt为空的边界情况：必须要处理无Prompt的“Everything”模式
现象：用户没有给点也没有给框，SAM应该输出全图所有分割掩码，但你的C++管线直接崩了。
原因：Prompt Encoder的输入为空时，需要传入一个默认的“背景点”坐标（比如[0,0]）以及对应的labels=[0]，否则TensorRT的动态shape为0时会异常。解决：在C++层对Prompt数量做保护，如果num_points == 0，初始化一个假的背景点，并把num_points置1。这是SAM官方的处理逻辑，别省略。
5.4 FP16下掩码边缘出现锯齿：不是Bug，是精度策略选择
现象：FP16引擎跑出来的掩码边缘有细碎的锯齿，尤其细长物体。
原因：FP16动态范围窄，LayerNorm和Softmax在FP16下的倒数误差放大到掩码边界。解决：如果对边缘质量要求高，可以只对Encoder的主干保持FP16，对Mask Decoder用FP32引擎。或者用TensorRT的Polygraphy做量化敏感层分析，把个别层强制回FP32。我一般直接对Decoder不做FP16，反正它计算量小。
5.5 多线程推理时CUDA context冲突，时延不稳定
现象：同一个引擎被多个线程调用，会出现随机增加10-20ms的峰值，甚至CUDA illegal memory access。
原因：IExecutionContext不是线程安全的，多线程共享同一个context会导致CUDA stream上的资源竞争。解决：每个线程创建独立的IExecutionContext实例，但共享ICudaEngine。引擎是不可变的，可以多线程读；context保存推理状态，必须隔离。这个坑在TCP服务接多个请求时必现。
6. 验证与进阶：用同一张图对齐Python结果，再谈多模型共享显存
做完C++推理，第一步不是跑Benchmark，而是“对结果”。拿一张作物分割图，用Python原版推理一次，把Mask数值序列化成npy，再在C++推理后也输出成npy，逐像素对比IoU。IoU高于0.95说明管线没问题，低于0.9就要检查预处理Resize是线性还是最近邻——SAM对插值方式极其敏感。对完结果再做性能验证：用trtexec测Engine的Throughput，再写一个C++计时函数，取1000帧的平均延迟，剔除前10次预热。我自己的经验是：端到端从Python的450ms降到C++的80ms，其中Encoder从380ms降到55ms，Decoder和前后处理占了剩下的25ms。这个量级说明TensorRT部署SAM分割模型C++代码和部署流程没白做。
再进一步，如果你要集成到YOLO检测流水线里，让SAM只在检测框内做分割，那么可以复用同一个TensorRT core，把YOLO的Engine和SAM的Engine放在同一个CUDA stream上。显存复用是关键：YOLO输出的是检测框，框坐标直接作为SAM的Prompt输入，两者共用一块显存buffer，用cudaMemcpyAsync在stream上串行执行，避免两个模型轮流加载卸载引擎。这里有个习惯：我每次部署都会写一个EngineManager类，统一管理所有引擎的加载、引用计数和显存峰值统计，否则两个模型叠加起来显存轻松超过4GB。另外，记得在C++里对SAM的输入图像做BGR转RGB——OpenCV默认是BGR，但SAM训练时用的是RGB，搞反了你的掩码会像负片一样。这条坑几乎每个新手都踩，我把它放在最后提醒你，希望帮到你。
本文还有配套的精品资源，点击获取
