简介面向需要将 SAM 模型部署到生产环境的 C 开发者这套资源以 TensorRT 为推理后端提供完整的加速落地实现。压缩包仅1.74MB共22个文件核心代码由7个头文件和2个C源文件组成涵盖模型加载、预处理、推理与后处理3个Notebook演示导出与TensorRT转换另有JSON配置、CMakeLists、README和Dockerfile支持跨平台构建。目前已有815人学习下载。资料内容覆盖从PyTorch导出、引擎转换、线程池并发到结果可视化并附有车辆图像与GIF动图展示分割效果。对希望在服务端或边缘设备获得实时细分能力的工程师这套源码与指南能有效减少TensorRT踩坑时间让SAM从Python原型顺利转为高性能C服务。1. 为什么SAM这种“分割一切”模型在C生产环境里跑不动“使用TensorRT部署SAM分割一切大模型C源码部署步骤.zip”这个标题我看第一眼就知道这又是一次从Python原型到C生产落地的硬仗。SAMSegment Anything Model是Meta开源的提示分割模型它最吸引人的地方是一个模型在任意图像上点什么就分割什么不需要为每个场景单独训练。但SAM的image encoder是ViT-H结构参数规模超过600M在PyTorch里跑一张图GPU推理也要一秒钟上下放到C服务里做实时交互分割根本扛不住。这个zip包解决的问题不是“怎么把SAM跑起来”而是“怎么把SAM跑到能用”。TensorRT是NVIDIA的高性能推理引擎能把ONNX模型做层融合、精度校准、kernel自动选择再编译成针对特定显卡优化的engine文件。C前端负责图像预处理、engine加载、执行推理、后处理这样整个链路从文件读取到结果输出都不经过Python解释器。这篇文章我会按完整链路拆先讲为什么用TensorRT而不是直接跑PyTorch再讲导出和构建engine的具体步骤然后给出一个能用的C推理骨架最后把部署中常见的坑集中列出来。适合的目标读者是已经在Python里跑通过SAM、准备把它弄到C服务或边缘设备上的人如果你是第一次接触TensorRT前三章也足够你理解全貌后面的代码可以直接抄着改。2. 部署前的四个选择模型结构、TensorRT版本、精度、动态shape2.1 不是所有模型都值得TensorRTSAM正好值得在做部署方案前我一般先判断这是个计算密集模型还是访存密集模型。SAM的image encoder是典型的计算密集ViT-H有几十层transformer块每次前向都在做大量矩阵乘法这正好是TensorRT最擅长优化的地方。而Mask Decoder参数量小、计算也少不值得单独折腾通常跟image encoder放同一个engine里。而像一些只做文本分类的BERT小模型TensorRT加成有限反而引入复杂度不划算。SAM这种“大encoder 小decoder”的结构优化重点全部集中在encoder上把FP16打开shape固定用TensorRT做kernel autotuning通常能比PyTorch快1.5到3倍显存占用还能更低。一个需要考虑的问题是SAM下载的原始权重是PyTorch格式TensorRT不认识需要一个中间渠道。常见方案是先用torch.onnx.export导出成ONNX再用trtexec或TensorRT的C API构建engine。这个链路在SAM上可行但有几个细节要注意我在第三章会具体写。2.2 TensorRT版本选择别追新要跟显卡和CUDA匹配TensorRT版本不是越新越好。新版TensorRT要求更新的CUDA版本或更新的显卡驱动如果你目标机器是Tesla T4这种老卡装最新版TensorRT反而可能因为驱动太旧跑不起来。我建议先nvidia-smi确认驱动版本和CUDA版本再倒推装哪个TensorRT。常见搭配是CUDA 11.8配TensorRT 8.5或8.6这套组合对T4、A30、A100都友好CUDA 12.0以上配TensorRT 8.6或9.x适合40系显卡和新款L20这种。装错了最典型的症状是运行时提示找不到libnvinfer.so.8或者报undefined symbol之类的链接错误。构建engine用的小工具trtexec通常跟TensorRT一起发布也支持直接用命令行转engine。但用trtexec只能验证模型有没有问题和跑benchmark生产环境最终还是要用C在程序里构建或加载engine。在源码包里开发者给的步骤文件里一般也会有环境准备说明按那个顺序来就行。2.3 FP16还是INT8SAM用FP16够用INT8要谨慎TensorRT的量化精度选择直接决定推理速度和精度损失。SAM的image encoder输出的是高维图像embedding后续要给decoder做相似度匹配精度非常敏感。我用FP16部署时分割结果肉眼基本看不出差异降到INT8后边缘细节会丢尤其对细小物体分割mask边界会出现毛刺。所以我的建议是第一版直接FP16跑通了再考虑INT8。INT8需要准备校准集对于SAM这种图像模型校准集要覆盖背景复杂、目标多变的图片如果校准集分布偏了精度损失会不可控这一块“玄学”成分很大除非你后续实测发现FP16延时仍不达标一般不建议第一版就上INT8。FP16在TensorRT里是个开关不用做任何额外校准。2.4 固定shape还是动态shapeSAM两者都要TensorRT的engine是对特定输入shape做优化的例如固定1080x1080x3的输入和固定批次1engine内部可以做更多层融合性能更好。但如果你的业务需要不同分辨率图片进来就得开动态shape用优化档位opt profile限制范围代价是推理延时会比固定shape高一些。SAM官方建议输入是1024x1024也就是说无论原图多大都要预处理好再丢给模型。如果业务场景就是一张张图单独推理输入分辨率固定1024那么固定shape是最好选择如果你要把多张不同尺寸的图拼batch一起推理就得上动态shape。源码包里如果作者已经帮你处理成固定shape那最好先按这个跑通后面再改。3. 从PyTorch权重到TensorRT engine导出、构建、验证全步骤3.1 先把torch模型导出成ONNX注意SAM的结构特殊性TensorRT不直接吃PyTorch权重这是部署SAM的第一个门槛。需要写一段Python脚本把sam的image_encoder和mask_decoder分别导出成两个ONNX文件或者合成一个。我实际做的时候发现合成一个文件更省事因为image encoder的输出embedding直接喂给mask decoder在同一个engine里做省的传递开销和显存拷贝都少很多。一个最基本导出脚本如下import torch from segment_anything import sam_model_registry sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth).cuda().eval() image torch.randn(1, 3, 1024, 1024).cuda() pt1 torch.tensor([[100.0, 100.0]], dtypetorch.float32).cuda() pt2 torch.tensor([[500.0, 500.0]], dtypetorch.float32).cuda() labels torch.tensor([[1, -1]], dtypetorch.int64).cuda() with torch.inference_mode(): # 先跑一次warmup确保所有模块被初始化 image_embeddings, _ sam.image_encoder(image) torch.onnx.export( sam.image_encoder, (image,), sam_image_encoder.onnx, opset_version17, input_names[image], output_names[image_embeddings], dynamic_axes{image: {0: batch}} )这段脚本里有两个关键点。一是opset_version我建议用17或以上TensorRT 8.6对opset17的支持比较稳。二是dynamic_axes这里只开了batch维度动态空间维度还是固定的1024x1024这样会让转换更稳定也防止某些op在动态shape下表现异常。导出后先别急着转engine用onnxruntime或onnxsimplifier过一遍能顺手去掉一堆无用节点。SAM的ViT-H结构里有大量层归一化LayerNorm和GeLU这类算子onnxsimplifier常常能合并一部分减小ONNX体积也让TensorRT解析时少一些弯路。3.2 用trtexec构建engine三种量化模式可以对比测导出好ONNX后最省事的构建方式是用trtexec命令行几个参数就能跑出结果。对于SAM这个模型我通常这么构建FP16 engine/opt/tensorrt/bin/trtexec \ --onnxsam_image_encoder.onnx \ --saveEnginesam_image_encoder.engine \ --fp16 \ --minShapesimage:1x3x1024x1024 \ --optShapesimage:1x3x1024x1024 \ --maxShapesimage:4x3x1024x1024如果不需要动态batch直接把minShapes和maxShapes去掉只保留--fp16。内存够的话还可以加一条--memPoolSizeworkspace:2048表示给TensorRT最多2GB的workspace空间做layer融合和kernel选择。给多了不一定更好但给少了经验上低于512MB会让TensorRT找不到更优kernel性能打折。构建完成后先用trtexec自带的benchmark跑一遍/opt/tensorrt/bin/trtexec --loadEnginesam_image_encoder.engine --shapesimage:1x3x1024x1024 --inputIOFormatsfp16看输出的FPS和显存占用能初步确认engine有没有生效FP16。如果延时不降反升基本可以肯定是某些层被强制回退到FP32或者shape范围开太大导致autotuning没有找到最优解。3.3 验证engine输出对不上的问题点engine构建成功不等于能用我至少踩过两次“构建成功但结果全错”的坑。基本验证方法是拿同一张图分别跑PyTorch和TensorRT计算二者分割mask的IoU或像素级差异。如果差异超过1%先怀疑精度问题如果完全对不上大概率是预处理或后处理不一致。这里要特别留意SAM的输入输出规范。image encoder的输入要求是RGB、像素范围0~255、resize到1024x1024并且要做归一化均值是[123.675, 116.28, 103.53]方差是[58.395, 57.12, 57.375]。很多人在Python里已经做对了但C重写预处理时忘了归一化顺序或像素范围导致embedding全偏了后面解码出的mask基本是无意义区域。另一个高频问题是mask decoder的输出SAM输出的是一个多通道的logits通常是3个mask的候选需要做sigmoid再根据iou_score选最优的那个。如果你只取了第一个通道恰好不是最优的效果就会差一点。这块属于后处理细节我在下一章的C代码里会直接把最优mask选择逻辑写出来。4. C推理管线源码骨架从加载engine到输出mask4.1 用C API构建engine最优做法是缓存成.plan文件反复加载整个部署源码里C部分的核心工作分为两大块engine的加载或构建以及推理执行。生产环境建议只在构建阶段使用TensorRT C API的builder构建完成后把engine序列化成.plan文件存到磁盘每次推理启动时直接用runtime反序列化加载省去重复构建的时间。SAM的ViT-H构建一次FP16 engine可能要几分钟到十几分钟每次重启服务都重新构建是灾难。下面是文件读取和runtime初始化的代码骨架#include NvInfer.h #include fstream #include vector std::vectorchar loadEngineFromFile(const std::string path) { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f.good()) { throw std::runtime_error(cannot open engine file: path); } size_t size f.tellg(); std::vectorchar blob(size); f.seekg(0, std::ios::beg); f.read(blob.data(), size); return blob; } nvinfer1::ICudaEngine* createEngine( nvinfer1::IRuntime* runtime, const std::vectorchar blob) { return runtime-deserializeCudaEngine(blob.data(), blob.size()); }每个TensorRT程序都必须创建runtime和context对推理来说才是真正干活的地方它负责管理推理时的中间激活值和临时显存。注意这个对象不是线程安全的多线程推理时每个线程都要维护自己的context不能共享。4.2 预处理C端做resize、归一化数据从CPU搬到GPU预处理在TensorRT里没有通用模板因为输入张量的布局、归一化参数每个模型不一样所以必须自己写。我见过最大的一个坑是resize算法不一致OpenCV的INTER_LINEAR和PyTorch的F.interpolate默认的bilinear算法对绝大多数图片结果接近但在纹理密集或对角线场景下会有几个像素的差异。SAM对图像非常敏感最稳妥的办法是统一用OpenCV的INTER_LINEAR进行resize因为SAM官方导出模型时就是用这类常规操作做预处理两者不会出现太大系统偏差。以下是预处理和输入buffer设置的代码骨架#include opencv2/opencv.hpp void preprocessImage(cv::Mat src, float* gpuInputPtr, cudaStream_t stream) { cv::Mat resized; cv::resize(src, resized, cv::Size(1024, 1024), 0, 0, cv::INTER_LINEAR); // BGR - RGB因为OpenCV默认读入是BGR cv::Mat rgb; cv::cvtColor(resized, rgb, cv::COLOR_BGR2RGB); // 归一化像素范围0~255原样保留到float下面再减均值除方差 cv::Mat floatImg; rgb.convertTo(floatImg, CV_32FC3); const float mean[3] {123.675f, 116.28f, 103.53f}; const float stddev[3] {58.395f, 57.12f, 57.375f}; float* ptr (float*)floatImg.data; std::vectorcv::Mat channels(3); cv::split(floatImg, channels); size_t planeSize 1024 * 1024; for (int c 0; c 3; c) { cv::Mat normalized; channels[c].convertTo(normalized, CV_32FC1, 1.0 / stddev[c], -mean[c] / stddev[c]); // 注意SAM权重是NHWC还是NCHW需要按engine输入格式决定 // 这里以NCHW为例每个channel先拷贝 } // 最后把三个channel拼到一块连续显存里再cudaMemcpyAsync到gpuInputPtr }这里最容易翻车的细节是“均值除方差”的运算顺序。有些代码把图像先乘1/255再减均值又是另一种归一化两者结果完全不同还很难一眼看出来。我的风格是只在预处理层做一次归一化模型内部本身不包含这个操作这样推理端更可控。同时尽量所有数据搬运都用cudaMemcpyAsync和推理同一根stream上避免CPU和GPU互等。4.3 推理执行申请显存buffer、绑定输入输出、executeV2engine确定后需要根据它的输入输出张量信息来分配显存buffer。有两种分配方式第一种是TensorRT推荐的IExecutionContext在每次推理时用setTensorAddress设置张量地址另一种是传统做法提前createBindings数组。新版API推荐前者但我自己为了稳定和兼容老显卡驱动一般就用executeV2加bindings数组维护成本低。一个基础的推理执行代码void runInference(nvinfer1::IExecutionContext* ctx, float* inputBuffer, float* outputEmbedding, int batchSize, cudaStream_t stream) { // 假设binding 0是inputbinding 1是output void* bindings[] { inputBuffer, outputEmbedding }; bool ok ctx-enqueueV2(bindings, stream, nullptr); if (!ok) { throw std::runtime_error(TensorRT inference failed); } }看代码很简单但需要注意一个前提——engine的输入shape和buffer大小必须匹配。如果你导出ONNX时batch是动态的C端要调用ctx-setInputShape(image, nvinfer1::Dims4{batchSize, 3, 1024, 1024})如果漏了这句话或者传进去的shape超出opt profile范围推理会直接报错。另外enqueueV2默认是异步的调用后结果不一定出来了必须在后续用cudaStreamSynchronize或cudaEvent来同步否则输出数据可能会读到脏数据。这是新手上手最常见的血泪经验之一。4.4 后处理从embedding生成mask多候选最优选择image encoder输出的embedding是1x256x64x64这种形状后续需要把它传给mask decoder。mask decoder吃三个输入image embedding、point prompt坐标以及point label表示正负样本点输出的是mask logits。C里如果想把整个SAM统一在一个engine里其实不需要自己写任何解码逻辑直接把它能过一遍TensorRT就行前提是导出ONNX的时候把整个SAM连在一起。如果你必须分两个engine来跑比如encoder和decoder是分开优化的那C端要把embedding先保存下来然后作为decoder输入。流程不复杂但需要额外的显存拷贝多一次H2D和D2H整体延时影响不明显。// 用sigmoid处理logits取最大的候选mask void postprocessMask(float* logits, int numCandidates, int H, int W, cv::Mat outputMask) { float bestScore -1e9f; int bestIdx 0; for (int c 0; c numCandidates; c) { float score logits[c]; if (score bestScore) { bestScore score; bestIdx c; } } // 从bestIdx对应位置拷贝做一个sigmoid再转为8位mask }这段代码只是一种示意实际SAM的mask decoder输出是1x3xHxW或者1x4xHxW前3个通道对应三种mask候选还有1个iou预测。惯例做法是三个候选mask分别sigmoid之后和原图大小还原到输入的尺寸再根据与提示点的关系决定选哪个。另外mask输出大小需要插值回用户看到的那张图这步别忘了。5. 部署避坑指南五个最容易翻车的地方5.1 坑一engine构建成功推理时输入shape不匹配直接报错这个现象是engine加载没问题调用enqueueV2时报“input shape invalid”或干脆崩掉。原因基本是导出时开了dynamic shape但C端没调setInputShape或者传的shape不在当时设置的min/opt/max范围内。解决方法是先确认你trtexec构建时的shape范围C端必须用同一套shape。建议第一版全部固定成1x3x1024x1024不做动态batch。5.2 坑二预处理归一化不一致mask结果完全对不上现象TensorRT跑出来的mask跟PyTorch结果差得离谱前景背景都颠倒了。常见原因是C端把图像先转成0~1的float再做减均值除方差而Python端是0~255直接做的。两者的结果差一个固定倍数关系后续二值化阈值全靠试怎么调都别扭。解决方法是写个单元测试固定喂一张纯色图对比Python和C预处理结果前几位float数值是否一致这样十有八九能定位问题。5.3 坑三多线程推理时context共享导致偶发性错误结果如果你用一个IExecutionContext在多个线程里并发调用TensorRT会随机给出错误结果甚至crash。原因是context内部维护了很多中间状态不是线程安全的。解决方法是每个线程一个contextengine可以共享但context必须一份一个。很多人忽略了这一点因为小batch测试时不会报错并发量一上去就“玄学”闪崩排查起来很耗时。5.4 坑四显存泄漏长时间运行后OOMTensorRT推理本身不会特别占显存但如果你每次推理new一个cudaMemcpy的buffer又忘了释放跑几万次后必爆显存。解决方法是启动时一次性分配所有输入输出buffer推理只复用不反复申请。这一点在C端格外突出Python有GC帮忙缓解C全靠自觉。5.5 坑五用trtexec测到的FPS和生产C代码测到的不一致trtexec默认做的是纯推理不包含预处理和后处理所以它报告的延时往往比实际端到端低很多。生产环境的延时必须自己测把预处理、H2D拷贝、推理、D2H拷贝、后处理全部计时累加这个数字才是真正的端到端性能。我见过有人拿trtexec的数字对外承诺P99延时上线直接翻车。所以你自己工程里要建立一套时间统计用cudaEvent记录GPU侧耗时用chrono记录CPU侧耗时两边都统计才能看到瓶颈在哪。6. 进阶用法把image embedding缓存复用做一个可交互的SAM服务SAM最实用的落地形态其实不是对每张图重新全量推理而是先对固定背景图像只做一次encoder推理缓存embedding接下来所有点击提示只需要执行一个很小的mask decoder。因为用户交互时图像不变只有点击的坐标在变image encoder结果可以复用这样交互延迟能做到几十毫秒级别。这个思路在源码包里如果没写我强烈建议你自己加。实现方法是把encoder和decoder拆成两个engine文件或者在程序里只运行一次encoder保存输出到显存后续每次click只跑decoder部分。用两个engine的方式更灵活但要注意decoder的输入里有坐标点数据需要CPU改成GPU端的float buffer。写C的时候把坐标点用cudaMalloc固定下来每次click前用cudaMemcpyAsync更新坐标点就行开销几乎可以忽略。这个方案落地后有个实际效果即便是SaaS服务在多用户共享GPU场景下第一个用户传来图片系统在500ms左右出基础结果后续每个点击只需几十ms就能回一个mask体验上已经“接近实时”。这就是TensorRT的一种架构价值。它还能做更多事情。比如固定背景做小目标检测之后再把小目标区域送去SAM分割就可以让检测和分割做联动。这在工业质检里非常常用先检测缺陷区域再对每个区域做精确分割最终输出缺陷面积、轮廓。这套链路里SAM作为后处理的精度担当TensorRT负责把整个流程压到实时以内。我在VinAI实践项目中遇到过一个典型的案例用SAM辅助半自动标注数据目标框的初始化是检测模型给的但如果边界不齐人工点击修正SAM加上TensorRT的高效交互一个人一天能精修1500张图比传统多边形框选快太多。正是因为交互延迟降下来了这种工作流才真正有人愿意用。最后想提醒一个习惯在整个C工程里建议你把所有TensorRT版本的兼容行为都写成一个独立的适配层。因为TensorRT 8.5和9.x的API接口有少量变动升级的时候改适配层就行不用动业务代码。我吃过这个亏用了半年后才把经验沉淀下来。希望这篇笔记帮到你少踩我踩过的那些“血泪坑”。本文还有配套的精品资源点击获取