简介这份资源面向具备一定 C 与深度学习部署基础的开发者聚焦于用 NVIDIA TensorRT 在 GPU 上高效推理 Segment Anything ModelSAM解决原始 PyTorch 版 SAM 推理速度慢、显存占用高、难以落地到实际工程的问题。项目名为 SPEED-SAM-C-TENSORRT通过 TensorRT 引擎构建与 CUDA 优化显著提升 GPU 利用率适合图像分割、交互式抠图、自动标注等场景的工程化部署。压缩包共 20 个文件约 71.22MB包含 3 个 cpp 源文件与多个 h 头文件构成的核心推理代码2 个 onnx 模型文件SAM 编码器与掩码解码器以及 jpg、png 示例图片和 txt、license 等辅助文件并配有 CMakeLists.txt 便于编译。已有 989 人学习下载。读者可据此获得一套完整的 C TensorRT 推理工程理解从 ONNX 模型到引擎序列化、CUDA 预处理与后处理的实现路径并借助示例图片快速验证分割效果为自研高性能视觉部署方案提供参考。1. 用 C 和 TensorRT 把 SAM 跑成生产力这条路到底值不值得走如果你手上有一个 SAMSegment Anything Model的.pth权重想把它塞进 C 服务里做实时分割大概率会卡在同一个地方Python 里三行代码就能出掩码搬到 C 里却要面对 ONNX 导出、TensorRT 引擎构建、显存管理、预处理对齐这一整套链路。标题里的「使用 C 中的 TensorRT 实现 SAM」说的就是把 Meta 这套分割模型从研究脚本变成可部署的推理服务这件事。它解决的不是「SAM 能不能用」而是「SAM 能不能在 C 工程里稳定、低延迟地用」。适合两类人一类是做视觉算法落地、需要把模型嵌进现有 C 框架的工程师另一类是想搞懂 TensorRT 部署全流程、拿 SAM 当练手项目的进阶开发者。这条路有门槛但走通之后你会发现所有 Transformer 类视觉模型的部署套路都是相通的。2. SAM 的 C TensorRT 部署链路从权重到引擎2.1 为什么不能直接把 .pth 丢给 TensorRTTensorRT 不认识 PyTorch 的权重格式它吃的是自己序列化出来的 plan 文件或者 ONNX 这类中间表示。所以第一步永远是格式转换。SAM 的结构比普通 CNN 复杂它有三个部分图像编码器Image Encoder基于 ViT、提示编码器Prompt Encoder、掩码解码器Mask Decoder。图像编码器是重计算的大头一张 1024×1024 的图跑一次 ViT-H 大概要几百毫秒提示编码器和掩码解码器很轻但它们是动态输入的——点、框、掩码提示的维度不固定。这就带来一个关键选型问题要不要把整个 SAM 导成一个 ONNX。我的建议是拆开导。图像编码器单独导成一个引擎输入固定为[1, 3, 1024, 1024]提示编码器和掩码解码器合并导成另一个引擎输入维度用动态 shape 处理。原因很直接图像编码器对同一张图只需要跑一次之后换不同的点提示可以复用它的输出 embedding。如果你把整个模型导成一个引擎每次换提示都要重跑 ViT延迟直接翻几倍。常见做法是先用 Python 导出 ONNX再用trtexec或 TensorRT 的 Python API 构建引擎。导出时注意 opset 版本SAM 里的grid_sample、einsum这类算子对 opset 有要求opset 17 是比较稳的选择。2.2 导出 ONNX 的两个关键脚本先看图像编码器的导出。下面这段代码在 Python 环境里跑依赖torch、segment-anything和onnximport torch from segment_anything import sam_model_registry # 加载官方权重vit_h 是精度最高但最重的版本 sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth) sam.eval() # 只取图像编码器包一层固定输入尺寸 class ImageEncoder(torch.nn.Module): def __init__(self, sam): super().__init__() self.encoder sam.image_encoder def forward(self, x): # 输出 shape: [1, 256, 64, 64] return self.encoder(x) encoder ImageEncoder(sam).cuda() dummy torch.randn(1, 3, 1024, 1024).cuda() torch.onnx.export( encoder, dummy, sam_image_encoder.onnx, input_names[image], output_names[embedding], opset_version17, do_constant_foldingTrue, )这段代码的逻辑是把image_encoder单独抽出来用一个固定尺寸的 dummy 输入触发 tracing。opset_version17是为了支持 ViT 里的 attention 算子。do_constant_foldingTrue会把能提前算的常量折叠掉减小 ONNX 体积。导出后你会得到一个几百 MB 的 ONNX 文件ViT-H 的编码器本身就大这是正常的。提示编码器和掩码解码器的导出稍微麻烦一点因为要处理动态输入。核心思路是把点坐标、点标签、框坐标都作为输入传进去import torch from segment_anything import sam_model_registry from segment_anything.modeling import PromptEncoder, MaskDecoder sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth) sam.eval() class PromptAndMask(torch.nn.Module): def __init__(self, sam): super().__init__() self.prompt_encoder sam.prompt_encoder self.mask_decoder sam.mask_decoder def forward(self, image_embedding, point_coords, point_labels): # sparse embeddings 来自点和框dense 这里先留空 sparse, dense self.prompt_encoder( points(point_coords, point_labels), boxesNone, masksNone, ) # 解码出掩码输出低分辨率 logits low_res_masks, _ self.mask_decoder( image_embeddingsimage_embedding, image_peself.prompt_encoder.get_dense_pe(), sparse_prompt_embeddingssparse, dense_prompt_embeddingsdense, multimask_outputFalse, ) return low_res_masks model PromptAndMask(sam).cuda() emb torch.randn(1, 256, 64, 64).cuda() coords torch.randn(1, 1, 2).cuda() # 一个点提示 labels torch.ones(1, 1, dtypetorch.int).cuda() torch.onnx.export( model, (emb, coords, labels), sam_prompt_mask.onnx, input_names[image_embedding, point_coords, point_labels], output_names[low_res_masks], opset_version17, dynamic_axes{ point_coords: {1: num_points}, point_labels: {1: num_points}, }, )这里dynamic_axes把点数量标成动态维度这样同一个引擎能处理 1 个点、5 个点甚至一个框框用两个点表示。multimask_outputFalse表示只输出一个掩码如果你需要 SAM 默认的三个候选掩码把它改成True输出维度会多一维。注意prompt_encoder里的get_dense_pe()依赖图像尺寸如果你的输入不是 1024这里要同步改。2.3 用 trtexec 构建引擎并验证ONNX 有了接下来构建 TensorRT 引擎。最省事的方式是trtexec它随 TensorRT 一起安装# 图像编码器固定 shape开 FP16 trtexec --onnxsam_image_encoder.onnx \ --saveEnginesam_image_encoder_fp16.plan \ --fp16 \ --workspace4096 # 提示解码器动态 shape需要指定 optimization profile trtexec --onnxsam_prompt_mask.onnx \ --saveEnginesam_prompt_mask_fp16.plan \ --fp16 \ --minShapespoint_coords:1x1x2,point_labels:1x1 \ --optShapespoint_coords:1x4x2,point_labels:1x4 \ --maxShapespoint_coords:1x16x2,point_labels:1x16 \ --workspace2048--fp16让 TensorRT 在支持 FP16 的显卡上用半精度计算速度通常能提升 30% 到 50%精度损失对分割任务几乎看不出来。--workspace是构建时允许用的显存上限单位 MBViT-H 建议给到 4096。动态引擎必须给minShapes、optShapes、maxShapes三个 profileoptShapes是你最常用的点数量TensorRT 会针对它做优化。如果你的点数量经常超过 16把maxShapes调大但注意显存占用会跟着涨。构建完成后用trtexec --loadEnginexxx.plan --shapes...可以快速验证引擎能不能跑通、延迟大概多少。这一步别跳过我见过太多人 ONNX 导出成功、引擎构建报错最后发现是某个算子不支持 FP16。3. C 侧推理代码把引擎跑起来3.1 TensorRT 运行时初始化与显存管理C 里用 TensorRT 的核心流程是反序列化 plan → 创建 execution context → 分配显存 → 绑定输入输出 → 执行。下面是一个最小可用的封装#include NvInfer.h #include cuda_runtime_api.h #include fstream #include vector #include memory class TrtEngine { public: TrtEngine(const std::string planPath) { // 读取 plan 文件到内存 std::ifstream file(planPath, std::ios::binary); file.seekg(0, std::ios::end); size_t size file.tellg(); file.seekg(0, std::ios::beg); std::vectorchar buffer(size); file.read(buffer.data(), size); // 创建 runtime 和 engine runtime_.reset(nvinfer1::createInferRuntime(logger_)); engine_.reset(runtime_-deserializeCudaEngine(buffer.data(), size)); context_.reset(engine_-createExecutionContext()); // 为每个输入输出分配显存 int n engine_-getNbBindings(); buffers_.resize(n); for (int i 0; i n; i) { auto dims engine_-getBindingDimensions(i); size_t bytes volume(dims) * sizeof(float); cudaMalloc(buffers_[i], bytes); bindingNames_[engine_-getBindingName(i)] i; } } ~TrtEngine() { for (auto ptr : buffers_) cudaFree(ptr); } void* buffer(const std::string name) { return buffers_[bindingNames_[name]]; } bool infer(cudaStream_t stream) { return context_-enqueueV2(buffers_.data(), stream, nullptr); } private: static size_t volume(const nvinfer1::Dims d) { size_t v 1; for (int i 0; i d.nbDims; i) v * d.d[i]; return v; } nvinfer1::Logger logger_; std::unique_ptrnvinfer1::IRuntime runtime_; std::unique_ptrnvinfer1::ICudaEngine engine_; std::unique_ptrnvinfer1::IExecutionContext context_; std::vectorvoid* buffers_; std::mapstd::string, int bindingNames_; };这段代码做了三件事反序列化引擎、创建执行上下文、按 binding 分配显存。volume函数把Dims转成元素个数乘以sizeof(float)得到字节数。注意这里假设所有输入输出都是 FP32如果你用了 FP16 引擎输入输出 binding 可能还是 FP32TensorRT 会自动转换但内部计算是 FP16。enqueueV2是异步执行配合 CUDA stream 可以做流水线。显存管理是 C 部署里最容易翻车的地方。SAM 的 ViT-H 编码器输出 embedding 是[1, 256, 64, 64]FP32 下就是 4MB不大但引擎内部的中间激活值可能占几百 MB。如果你同时跑多个实例显存会爆。我的习惯是给每个引擎单独一个 context不要共享因为 context 之间切换有开销。3.2 预处理与后处理对齐 Python 的数值C 侧最容易出错的不是推理本身而是预处理。Python 里 SAM 的预处理是resize 到 1024×1024保持长宽比短边补零、归一化减 mean 除 std、转 CHW。C 里如果直接用 OpenCV 的resize插值方式和 PyTorch 不一致会导致 embedding 有细微偏差最终掩码边缘对不上。#include opencv2/opencv.hpp // 输入 BGR 图输出归一化后的 NCHW float 数据 void preprocess(const cv::Mat bgr, float* dst, int targetSize 1024) { int h bgr.rows, w bgr.cols; float scale static_castfloat(targetSize) / std::max(h, w); int nh static_castint(h * scale); int nw static_castint(w * scale); cv::Mat resized; // 用 INTER_LINEAR 和 PyTorch 的 bilinear 对齐 cv::resize(bgr, resized, cv::Size(nw, nh), 0, 0, cv::INTER_LINEAR); // 放到 1024x1024 画布上右下补零 cv::Mat canvas cv::Mat::zeros(targetSize, targetSize, CV_8UC3); resized.copyTo(canvas(cv::Rect(0, 0, nw, nh))); // BGR - RGB, 归一化, HWC - CHW cv::Mat rgb; cv::cvtColor(canvas, rgb, cv::COLOR_BGR2RGB); rgb.convertTo(rgb, CV_32FC3, 1.0 / 255.0); std::vectorcv::Mat channels(3); cv::split(rgb, channels); float mean[3] {0.485f, 0.456f, 0.406f}; float stdv[3] {0.229f, 0.224f, 0.225f}; for (int c 0; c 3; c) { channels[c] (channels[c] - mean[c]) / stdv[c]; memcpy(dst c * targetSize * targetSize, channels[c].data, targetSize * targetSize * sizeof(float)); } }关键点是INTER_LINEAR和 PyTorch 的bilinear对齐以及归一化参数必须和训练时一致。mean和std是 ImageNet 的标准值SAM 用的就是这套。补零的位置也要注意SAM 官方是右下补零如果你补在左上embedding 会偏。后处理是把low_res_masks通常是 256×256上采样回原图尺寸然后二值化。上采样同样用双线性阈值一般取 0。如果你需要更精细的边缘可以在原图分辨率上做一次条件随机场但那是另一个话题了。3.3 一个完整的单点提示推理流程把上面的片段串起来一次完整的「给一个点、出一张掩码」流程是这样的void segmentByPoint(const cv::Mat image, float px, float py) { TrtEngine encoder(sam_image_encoder_fp16.plan); TrtEngine decoder(sam_prompt_mask_fp16.plan); // 1. 预处理图像 std::vectorfloat input(3 * 1024 * 1024); preprocess(image, input.data()); // 2. 拷贝到显存跑图像编码器 cudaMemcpy(encoder.buffer(image), input.data(), 3 * 1024 * 1024 * sizeof(float), cudaMemcpyHostToDevice); encoder.infer(0); // 3. 准备点提示坐标要归一化到 [0,1] float coords[2] {px / image.cols, py / image.rows}; int labels[1] {1}; // 1 表示前景点 cudaMemcpy(decoder.buffer(point_coords), coords, 2 * sizeof(float), cudaMemcpyHostToDevice); cudaMemcpy(decoder.buffer(point_labels), labels, sizeof(int), cudaMemcpyHostToDevice); // 4. 把编码器的输出 embedding 拷到解码器输入 cudaMemcpy(decoder.buffer(image_embedding), encoder.buffer(embedding), 256 * 64 * 64 * sizeof(float), cudaMemcpyDeviceToDevice); // 5. 跑解码器拿回掩码 decoder.infer(0); std::vectorfloat masks(256 * 256); cudaMemcpy(masks.data(), decoder.buffer(low_res_masks), 256 * 256 * sizeof(float), cudaMemcpyDeviceToHost); // 6. 上采样回原图并二值化 cv::Mat maskLow(256, 256, CV_32F, masks.data()); cv::Mat maskFull; cv::resize(maskLow, maskFull, image.size(), 0, 0, cv::INTER_LINEAR); cv::Mat binary maskFull 0.0f; }这段代码里点坐标归一化到[0,1]是 SAM 的要求不是像素坐标。labels里 1 是前景点0 是背景点如果你想排除某个区域可以加一个 label 为 0 的点。cudaMemcpyDeviceToDevice那一步是把编码器输出直接拷到解码器输入避免了绕回主机内存这是性能优化的关键。实际工程里编码器只需要对每张图跑一次之后换点提示只跑解码器所以你会把编码器的输出缓存起来。4. 避坑与排查SAM TensorRT 部署的五个血泪教训4.1 引擎构建报「unsupported operator」现象trtexec构建时报某个算子不支持常见的是GridSample或Einsum。原因TensorRT 版本和 ONNX opset 不匹配或者该算子在你的 TensorRT 版本里没有 FP16 实现。解决先确认 TensorRT 版本8.5 以上对GridSample支持较好如果 FP16 不支持去掉--fp16用 FP32 构建或者用--layerPrecisions单独指定某些层用 FP32。4.2 掩码边缘和 Python 结果对不上现象C 出的掩码比 Python 小一圈或者边缘锯齿明显。原因预处理 resize 的插值方式不一致或者补零位置不同。解决把 C 预处理后的图存成 npy和 Python 的预处理结果逐像素对比差异应该小于 1e-3。重点检查INTER_LINEAR和补零的Rect起点。4.3 动态 shape 引擎第一次推理特别慢现象动态引擎第一次跑某个点数量时延迟很高第二次就正常了。原因TensorRT 对动态 shape 需要在实际 shape 上做一次 tactic 选择第一次是 warmup。解决在服务启动时用optShapes对应的点数量跑几次 warmup把 tactic 缓存住。别等到线上第一个请求来了才 warmup。4.4 显存泄漏导致跑几十次后 OOM现象服务跑一段时间后cudaMalloc失败。原因每次推理都新建 context 或忘记释放 buffer。解决context 和 buffer 在初始化时创建一次复用如果必须动态创建确保析构里cudaFree。用nvidia-smi观察显存曲线正常应该是平的。4.5 FP16 引擎精度下降导致小目标漏分割现象大目标分割正常小目标掩码缺失。原因FP16 的动态范围有限小目标的 logits 值很小被截断成 0。解决对掩码解码器保持 FP32只对图像编码器用 FP16。或者用--fp16构建后在 C 里对输出做一次阈值补偿但更稳的做法是解码器用 FP32。5. 进阶技巧让 SAM 在 C 里跑得更快更稳5.1 用 CUDA Graph 消除 kernel 启动开销当你的点提示数量固定时整个推理流程的 kernel 序列是确定的这时候可以用 CUDA Graph 把一串 kernel 捕获成一个图一次提交。对于解码器这种小模型kernel 启动开销可能占延迟的一半。做法是在 warmup 阶段用cudaStreamBeginCapture和cudaStreamEndCapture把enqueueV2包起来之后每次推理直接cudaGraphLaunch。注意 TensorRT 的enqueueV2在 capture 模式下要用enqueueV3或者确保没有同步操作。5.2 批量处理多个点提示如果你的场景是「一张图、多个点、要多个掩码」别循环调用解码器。把点提示拼成 batch比如 8 个点提示拼成[8, 1, 2]的 coords一次推理出 8 个掩码。解码器很轻batch 8 的延迟比 batch 1 高不了多少但吞吐量翻 8 倍。前提是导出 ONNX 时把 batch 维度也标成动态。5.3 引擎版本兼容性检查TensorRT 的 plan 文件不是向前兼容的8.x 构建的引擎在 10.x 上可能加载失败。如果你的部署环境 TensorRT 版本会变要么每次重新构建引擎要么在代码里做版本检查int32_t major, minor, patch; runtime_-getEngineVersion(major, minor, patch); // 和构建时的版本对比不一致就重新构建我一般会在服务启动时打印引擎的版本信息出问题时第一眼就能看到是不是版本不匹配。5.4 一个我踩过的坑别在析构里做同步早期我写的封装在析构函数里调用了cudaStreamSynchronize结果服务退出时偶尔卡死。原因是析构顺序不确定stream 可能已经被销毁了。后来改成显式调用shutdown()方法在里面做同步和释放析构只做兜底。这个习惯帮我省了很多「玄学」崩溃。部署 SAM 这类模型最深的体会是Python 里一行predictor.predict()背后是 C 里几十个显存拷贝、shape 对齐和版本检查。但一旦跑通你会发现换任何 ViT 类模型都是同一套骨架。我现在的习惯是每接一个新模型先花半天把 ONNX 导出和引擎构建的脚本固化下来后面调 C 就是填空。希望帮到你。本文还有配套的精品资源点击获取