简介面向需要将 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快一个量级做图像分割的同行应该都有这种感觉SAMSegment 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)就够了batch1的实时场景没人会用batch8跑分割。还有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走TRTDecoder用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](checkpointsam_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_version17, input_names[image], output_names[image_embedding], dynamic_axesNone, # 固定shape避免动态轴 )这里必须设置dynamic_axesNone否则导出会生成动态shape的ONNX后面TRT转换时虽然能转但会禁用很多优化速度反而不如静态。opset_version建议用17或18高版本对LayerNorm、Gelu的支持更完整。别忘了torch.onnx.export默认会走traceSAM里的插值操作和mask_decode无关但Encoder内部有window_position这种张量操作固定输入尺寸100%避免trace出错。导出后用onnxsim精简一下模型能干掉一部分冗余Reshape和Transpose。3.2 用trtexec生成FP16引擎先跑通再写代码TRT引擎生成最快的方式是用trtexec命令行不用先写C。命令如下/opt/TensorRT/bin/trtexec \ --onnxsam_encoder.onnx \ --saveEnginesam_encoder_fp16.engine \ --fp16 \ --minShapesimage:1x3x1024x1024 \ --optShapesimage:1x3x1024x1024 \ --maxShapesimage:1x3x1024x1024 \ --workspace4096参数说明--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::vectorchar blob(size); file.read(blob.data(), size); // 创建runtime和engine runtime_ std::unique_ptrnvinfer1::IRuntime(nvinfer1::createInferRuntime(logger_)); engine_ std::unique_ptrnvinfer1::ICudaEngine(runtime_-deserializeCudaEngine(blob.data(), size)); context_ std::unique_ptrnvinfer1::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_ptrnvinfer1::IRuntime runtime_; std::unique_ptrnvinfer1::ICudaEngine engine_; std::unique_ptrnvinfer1::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::vectorfloat 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.2padding到1024×1024点坐标要乘scale后加上左上偏移。这个换算错了Prompt就点不到物体上。Mask Decoder的输入除了Prompt embedding外还需要Image Encoder的embedding所以C管线里必须保持embedding的显存驻留不能每次从GPU拷回CPU再传回去。我建议在GPU上分配embedding_bufferEncoder输出直接留在显存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_gridbilinear sampling组合但工作量大。我实际用的办法是用opset_version17TensorRT 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降到55msDecoder和前后处理占了剩下的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搞反了你的掩码会像负片一样。这条坑几乎每个新手都踩我把它放在最后提醒你希望帮到你。本文还有配套的精品资源点击获取