SAM模型TensorRT C++部署实战：从ONNX到低延迟推理

简介本资源是一套基于TensorRT加速的SAMSegment Anything Model大模型C部署完整实现面向具备C与深度学习推理基础的算法工程师、嵌入式AI开发者及高性能计算实践者解决视觉分割模型在端侧或服务端高效落地的核心难题。压缩包共22个文件涵盖2个核心CPP主程序、7个头文件含SAM核心算子封装与TensorRT引擎构建逻辑、3个Jupyter Notebook含模型导出与验证流程、2份Markdown文档含Windows/Linux双平台部署指南与中文说明、以及GIF动图演示、JPG/PNG测试图像和Docker开发环境配置文件整体仅1.74MB轻量但结构完整。已有814人学习下载资源提供从ONNX模型转换、TensorRT序列化引擎构建、内存缓冲管理到多线程推理调用的全链路C源码目录模块清晰含ThreadPool.h等工程级组件并附带VIM-H变体模型导出教程与实测效果截图可直接编译运行并快速集成至现有C项目。1. 为什么用 TensorRT 部署 SAM 大模型必须写 C不是 Python 不够快而是推理链路里藏着三个「不可绕过」的硬约束SAMSegment Anything Model在边缘端或工业级服务中落地时常卡在同一个地方Python 推理看似跑通了但一上产线就抖——GPU 显存暴涨、首帧延迟超 800ms、多路并发直接 OOM。这不是模型本身的问题而是 PyTorch 的 Python runtime 在部署侧天然带三重枷锁GIL 锁死多线程吞吐、动态图执行无法做 kernel 级融合、Tensor 内存生命周期由 GC 控制导致显存碎片化。而 TensorRT C 的组合是目前唯一能把 SAM 的 encoder-decoder 结构压进 120ms 内、显存占用稳定在 1.8GB 以下、且支持 RK3588 / Jetson Orin / A100 多平台复用的路径。本方案不依赖任何 Python 运行时所有 tensor 生命周期、stream 同步、内存池分配全由 C 手动控制zip 包里提供的sam_trt_engine.cpp和trt_sam_inference.h是经过 7 轮实测打磨的最小可行封装不是 demo是能直接塞进安防摄像头固件、工业质检 SDK 或车载视觉中间件里的生产级代码。适合正在把 SAM 接入 C 主程序、需要低延迟高吞吐、且不愿被 Python 环境绑架的嵌入式/算法/部署工程师。2. 从 .pt 到 .engineSAM 模型 TensorRT 引擎构建的四步闭环SAM 官方发布的sam_vit_h.pth是 PyTorch 格式不能直接喂给 TensorRT。必须走一条「导出 → 优化 → 序列化 → 加载」的确定性路径。关键不是“能不能转”而是“转出来的 engine 是否保留了 prompt embedding 的动态 shape 支持”——这是 SAM 实现“任意点/框分割”的核心能力也是多数失败案例的根源。2.1 导出 ONNX冻结 backbone 动态输入声明SAM 的 encoderViT-H是静态结构但 decoder 接收的 prompt embeddingpoint/box/mask是动态 batch 和 variable length。ONNX 导出时若未显式声明 dynamic_axes后续 TensorRT 将拒绝构建可变输入引擎import torch import onnx from segment_anything import sam_model_registry, SamPredictor # 加载原始模型注意必须用官方 checkpoint非 finetune 版 sam sam_model_registry[vit_h](checkpointsam_vit_h.pth) sam.eval() # 构造 dummy inputencoder 输入固定为 (1,3,1024,1024)prompt 保持动态 dummy_image torch.randn(1, 3, 1024, 1024) dummy_points torch.randn(1, 1, 2) # [B, N, 2] dummy_labels torch.randint(0, 2, (1, 1)) # [B, N] # 关键指定 dynamic_axes让 points/labels 的 N 维可变 torch.onnx.export( sam.image_encoder, dummy_image, sam_encoder.onnx, opset_version17, input_names[image], output_names[image_embedding], dynamic_axes{image: {0: batch, 2: height, 3: width}} ) # decoder 单独导出需 patch官方 model.forward_mask_decoder 未暴露为 callable # 实际 zip 包中已提供 patched_decoder.py此处省略 patch 细节 torch.onnx.export( patched_decoder, (dummy_image, dummy_points, dummy_labels, torch.zeros(1, 256, 64, 64)), sam_decoder.onnx, opset_version17, input_names[image, points, labels, image_embedding], output_names[masks, iou_predictions], dynamic_axes{ points: {1: num_points}, labels: {1: num_points}, image_embedding: {0: batch} } )逻辑说明SAM 的 encoder 只处理图像输出固定 shape 的 image_embeddingdecoder 才真正消耗 prompt。因此必须分两段导出且 decoder 的points和labels的第二维点数必须设为 dynamic_axes否则 TensorRT 编译时会报Unsupported shape inference。参数说明opset_version17是底线——低于此版本不支持aten::where等 SAM decoder 中大量使用的算子dynamic_axes的 key 必须与input_names严格一致否则 TensorRT 加载 ONNX 时无法映射 shape。2.2 TensorRT Builder 配置启用 FP16 优化 profile 显存策略ONNX 只是中间表示真正决定性能的是 TensorRT builder 的配置。SAM 的 ViT-H 参数量大632M但计算密集度远低于 LLM因此重点不在 INT8 量化会损 mask 边缘精度而在 kernel 融合和 memory pool 分配// trt_builder_config.cppzip 包中已实现 IBuilder* builder createInferBuilder(logger); INetworkDefinition* network builder-createNetworkV2(0); // 解析 ONNX注意必须用 parser-parseFromFile而非 parse否则 dynamic shape 丢失 auto parser nvonnxparser::createParser(*network, logger); parser-parseFromFile(sam_encoder.onnx, static_castint(ILogger::Severity::kWARNING)); // 关键配置启用 FP16但禁用 INT8SAM 对数值敏感 builder-setFP16Mode(true); builder-setInt8Mode(false); // 不要强行 INT8 // 设置 optimization profileSAM 输入图像固定为 1024x1024但 prompt 点数可变 IOptimizationProfile* profile builder-createOptimizationProfile(); profile-setDimensions(image, OptProfileSelector::kMIN, Dims4{1,3,1024,1024}); profile-setDimensions(image, OptProfileSelector::kOPT, Dims4{1,3,1024,1024}); profile-setDimensions(image, OptProfileSelector::kMAX, Dims4{1,3,1024,1024}); // prompt 输入维度必须覆盖典型范围 profile-setDimensions(points, OptProfileSelector::kMIN, Dims3{1,1,2}); profile-setDimensions(points, OptProfileSelector::kOPT, Dims3{1,3,2}); // OPT 设为 3 点兼顾速度与灵活性 profile-setDimensions(points, OptProfileSelector::kMAX, Dims3{1,16,2}); // MAX 设为 16 点防极端 case builder-addOptimizationProfile(profile); // 显存策略SAM encoder 输出 embedding 约 256MB必须预分配足够 workspace builder-setMaxWorkspaceSize(2ULL 30); // 2GB workspace低于此值编译失败率超 60%逻辑说明OptimizationProfile不是可选配置而是 SAM 动态 prompt 的强制要求。TensorRT 需要知道points维度的 min/opt/max 三元组才能生成对应 kernel若只设 OPT运行时遇到 10 个点就会 crash。参数说明setMaxWorkspaceSize(2ULL 30)中的2ULL表示 unsigned long long避免 int 溢出 30是 2GB 的位运算写法比2 * 1024 * 1024 * 1024更安全。workspace 不足会导致 builder 返回 nullptr且无明确 error message——这是新手最常踩的静默坑。2.3 序列化 engine 并保存为 plan 文件编译完成的 engine 是内存对象必须序列化为.plan文件才能脱离 build 环境加载IHostMemory* serialized_engine engine-serialize(); std::ofstream p(sam_vit_h.engine, std::ios::binary); p.write(reinterpret_castconst char*(serialized_engine-data()), serialized_engine-size()); p.close(); serialized_engine-destroy();逻辑说明.plan是 TensorRT 的二进制序列化格式包含 optimized kernel、memory layout 和 binding 信息。它与 build 时的 TensorRT 版本、CUDA 版本、GPU 架构强绑定——A100 上生成的 plan 不能直接在 RTX 4090 上运行必须重新 build。参数说明serialized_engine-size()返回字节数必须用reinterpret_castconst char*强转否则 write 会截断。zip 包中的build_engine.sh已封装此流程并自动校验 CUDA_ARCHITECTURE如-gencode archcompute_80,codesm_80for A100。3. C 推理引擎封装如何让 SAM 在 120 行内完成一次完整分割zip 包中的trt_sam_inference.h是核心——它把 TensorRT 的 raw API 封装成SamTrtSession类屏蔽 device memory 分配、stream 同步、binding 绑定等细节开发者只需关注输入/输出 tensor 的 CPU-GPU 拷贝。3.1 初始化加载 engine 创建 execution context// trt_sam_inference.h 中的构造函数关键片段 SamTrtSession(const std::string engine_path) { // 1. 读取 .plan 文件 std::ifstream file(engine_path, std::ios::binary | std::ios::ate); std::streamsize size file.tellg(); file.seekg(0, std::ios::beg); std::vectorchar buffer(size); file.read(buffer.data(), size); // 2. 反序列化 engine runtime_ createInferRuntime(logger_); engine_ runtime_-deserializeCudaEngine(buffer.data(), size, nullptr); context_ engine_-createExecutionContext(); // 3. 获取 binding 数量SAM encoderdecoder 共 6 个 binding num_bindings_ engine_-getNbBindings(); bindings_ std::vectorvoid*(num_bindings_); // 4. 分配 device memory注意必须按 binding index 顺序 malloc for (int i 0; i num_bindings_; i) { Dims dims engine_-getBindingDimensions(i); size_t volume 1; for (int j 0; j dims.nbDims; j) volume * dims.d[j]; size_t bytes volume * sizeof(float); cudaMalloc(bindings_[i], bytes); } }逻辑说明bindings_是 void* 数组索引 0~2 对应 encoder 输入/输出3~5 对应 decoder 输入/输出。TensorRT 不保证 binding 顺序与 ONNX name 一致必须用engine_-getBindingName(i)检查zip 包中check_bindings.cpp提供验证脚本。参数说明cudaMalloc分配的显存大小由volume * sizeof(float)计算SAM 的 image_embedding 是float32不能误用sizeof(half)否则 decoder 输入错位导致 mask 全黑。3.2 执行分割从 CPU 图像到 GPU mask 的七步数据流std::vectorcv::Mat SamTrtSession::segment(const cv::Mat image_bgr, const std::vectorcv::Point points, const std::vectorint labels) { // 步骤 1BGR→RGB resize to 1024x1024 normalize cv::Mat image_rgb, image_resized; cv::cvtColor(image_bgr, image_rgb, cv::COLOR_BGR2RGB); cv::resize(image_rgb, image_resized, cv::Size(1024, 1024)); float* h_image new float[1024*1024*3]; // ... 归一化到 [0,1] 并 HWC→CHW // 步骤 2拷贝图像到 GPU cudaMemcpy(bindings_[0], h_image, 1024*1024*3*sizeof(float), cudaMemcpyHostToDevice); // 步骤 3构造 prompt tensorpoints/labels float* h_points new float[points.size()*2]; int* h_labels new int[labels.size()]; // ... 填充数据 cudaMemcpy(bindings_[1], h_points, points.size()*2*sizeof(float), cudaMemcpyHostToDevice); cudaMemcpy(bindings_[2], h_labels, labels.size()*sizeof(int), cudaMemcpyHostToDevice); // 步骤 4执行 encoder获取 image_embedding context_-setBindingDimensions(0, Dims4{1,3,1024,1024}); context_-setBindingDimensions(1, Dims3{1,(int)points.size(),2}); context_-setBindingDimensions(2, Dims2{1,(int)labels.size()}); context_-executeV2(bindings_.data()); // 步骤 5从 GPU 读回 embedding用于 decoder 输入 float* d_embedding new float[256*64*64]; cudaMemcpy(d_embedding, bindings_[2], 256*64*64*sizeof(float), cudaMemcpyDeviceToHost); // 步骤 6设置 decoder binding 并执行 cudaMemcpy(bindings_[3], d_embedding, 256*64*64*sizeof(float), cudaMemcpyHostToDevice); context_-setBindingDimensions(3, Dims4{1,256,64,64}); context_-executeV2(bindings_.data()); // 步骤 7读取 masksshape: [1,3,256,256]并 resize 回原图尺寸 float* h_masks new float[3*256*256]; cudaMemcpy(h_masks, bindings_[5], 3*256*256*sizeof(float), cudaMemcpyDeviceToHost); // ... 后处理argmax 取最佳 maskbilinear resizecv::Mat 包装 return result_masks; }逻辑说明executeV2是 TensorRT 8.4 的推荐接口替代已废弃的execute必须在每次 execute 前调用setBindingDimensions因为 SAM 的 prompt 点数可变TensorRT 需要 runtime 重规划 kernel。参数说明bindings_[5]是 decoder 的输出 maskshape 为[1,3,256,256]3 个 mask proposal不是[1,1,256,256]——SAM 总是输出 top-3 masks需取iou_predictions最大的那个zip 包中postprocess.cpp实现。4. 避坑SAM TensorRT C 部署的五个血泪经验现象、原因、解决不讲虚的全是实测翻车现场。4.1 现象context_-executeV2()返回 false但getLastError()为空字符串原因TensorRT 未正确设置 optimization profile 的 dimensions尤其points维度的 min/opt/max 三元组缺失或超出实际输入范围。例如 profile 中pointsmax 设为 10但传入 12 个点TensorRT 不报错也不执行静默失败。解决在executeV2前加断言检查auto dims context_-getBindingDimensions(1); // points binding if (dims.d[1] points.size()) { throw std::runtime_error(Points count exceeds TRT profile max); }4.2 现象mask 边缘严重锯齿IoU 低于 0.6官方 benchmark 为 0.8原因ONNX 导出时未冻结 batch norm或 TensorRT builder 未启用setFP16Mode(true)导致 float32 计算引入累积误差。ViT 的 layer norm 对数值 scale 敏感。解决导出前对 model 调用model.eval().train(False)并在 builder 配置中强制setFP16Mode(true)若仍不准改用setPrecisionConstraints(PrecisionConstraint::kSTRICT)。4.3 现象多线程调用segment()时偶发 segmentation fault原因context_-executeV2()不是线程安全的多个线程共用同一 context 会竞争 stream 和 memory。解决为每个线程创建独立ExecutionContextengine_-createExecutionContext()或使用context_-setOptimizationProfileAsync(0, stream)绑定专属 CUDA stream。4.4 现象RK3588 上加载 engine 失败报Could not find plugin: ResizeNearest_TRT version 1原因ONNX 中的Resize算子被解析为 TensorRT 插件但 RK3588 的 TensorRT 版本8.5.2.2不包含该插件需手动注册。解决在createInferRuntime后添加initLibNvInferPlugins(logger_, ); // 并链接 -lnvinfer_plugin4.5 现象cudaMalloc失败错误码 2out of memory原因setMaxWorkspaceSize设得太小或未释放之前cudaMalloc的显存C 忘记 delete[] 导致显存泄漏。解决用nvidia-smi监控显存在SamTrtSession析构函数中显式cudaFree所有 bindingsworkspace 至少设为 1.5GBSAM ViT-H 最小需求。5. 进阶技巧如何把 SAM 的 prompt embedding 预计算做到 15ms 内SAM 的瓶颈不在 decoder而在 encoder——ViT-H 对 1024x1024 图像的 forward 需要 45~60msA100。但实际业务中同一张图常被多次分割不同点/框重复跑 encoder 是巨大浪费。zip 包中的cache_manager.h实现了基于图像 hash 的 embedding 缓存把二次分割的 encoder 时间压到 15ms 内。5.1 图像指纹生成不用 full MD5用 perceptual hashMD5 对 resize 后的图像不鲁棒而感知哈希phash能容忍几何变换和轻微压缩uint64_t compute_phash(const cv::Mat image) { cv::Mat gray, resized; cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY); cv::resize(gray, resized, cv::Size(32, 32)); cv::Mat float_img; resized.convertScaleAbs(float_img, 1.0/255.0); cv::Mat dct; cv::dct(float_img, dct); cv::Mat top_left dct(cv::Rect(0,0,8,8)); double mean_val cv::mean(top_left)[0]; uint64_t hash 0; for (int i 0; i 8; i) { for (int j 0; j 8; j) { if (dct.atfloat(i,j) mean_val) hash | (1ULL (i*8j)); } } return hash; }逻辑说明8x8 DCT 系数矩阵的均值作为阈值生成 64-bit 指纹。实测对同一张图 resize 后 phash 相同率 99.9%而 MD5 相同率为 0%。参数说明1ULL (i*8j)使用 unsigned long long 避免左移溢出cv::dct输入必须是CV_32F需提前convertScaleAbs归一化。5.2 LRU cache with CUDA memory pinning缓存必须支持快速查找 自动淘汰且 embedding 显存不能频繁 malloc/freeclass EmbeddingCache { private: std::unordered_mapuint64_t, cudaStream_t streams_; std::unordered_mapuint64_t, void* embeddings_; std::listuint64_t lru_order_; const size_t max_size_ 100; // 最多缓存 100 张图 public: void put(uint64_t hash, void* embedding_ptr) { if (embeddings_.find(hash) embeddings_.end()) { if (embeddings_.size() max_size_) { // 淘汰 LRU 最久未用项 uint64_t oldest lru_order_.back(); cudaFree(embeddings_[oldest]); embeddings_.erase(oldest); lru_order_.pop_back(); } } // 插入新项到 front lru_order_.push_front(hash); embeddings_[hash] embedding_ptr; } void* get(uint64_t hash) { auto it embeddings_.find(hash); if (it ! embeddings_.end()) { // 移动到 frontMRU lru_order_.erase(std::find(lru_order_.begin(), lru_order_.end(), hash)); lru_order_.push_front(hash); return it-second; } return nullptr; } };逻辑说明cudaFree必须在put中显式调用否则显存泄漏lru_order_用std::list而非std::vector避免 erase 时迭代器失效。参数说明max_size_100是经验值——每张图 embedding 占 2566464*44MB100 张约 400MB 显存远低于 A100 的 40GB安全水位。5.3 预热策略冷启动时主动填充 cache首次调用segment()前用空图像触发 encoder warmup避免首帧抖动void SamTrtSession::warmup() { cv::Mat dummy(1024,1024,CV_8UC3,cv::Scalar(128,128,128)); std::vectorcv::Point dummy_pts {{512,512}}; std::vectorint dummy_lbl {1}; segment(dummy, dummy_pts, dummy_lbl); // 执行一次触发 kernel 加载 }我上线过 3 个 SAM 边缘项目每次都在warmup()这行代码上栽过跟头——没它首帧延迟 1200ms有它稳定在 110ms。这行代码就是我的后悔药。希望帮到你。本文还有配套的精品资源点击获取