ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

SAM模型TensorRT C++部署实战:从PyTorch权重到高性能推理引擎

SAM模型TensorRT C++部署实战:从PyTorch权重到高性能推理引擎 简介这份资源面向具备一定 C 与深度学习部署基础的开发者聚焦于用 NVIDIA TensorRT 高效推理 Segment Anything ModelSAM解决原始 PyTorch 版 SAM 在 GPU 上推理速度慢、显存占用高的问题。项目名为 SPEED-SAM-C-TENSORRT通过 TensorRT 引擎构建与 CUDA 优化显著提升 GPU 利用率适合图像分割、交互式抠图、边缘端视觉部署等场景。压缩包共 20 个文件约 71.22MB包含 3 个 cpp 源文件与 7 个头文件构成的核心推理代码、2 个 onnx 模型文件SAM 编码器与掩码解码器、CMakeLists 构建脚本以及若干 jpg、png 示例图片和 txt、license 说明文件目录按头文件、源码、模型、资源分层组织。目前已有 989 人学习下载。读者可据此获得一套可直接编译运行的 C TensorRT 推理工程理解 SAM 编码器与掩码解码器的引擎封装、CUDA 内存管理及前后处理流程并借助示例图片快速验证分割效果为二次开发与性能调优提供参考。1. 从 PyTorch 权重到 TensorRT 引擎SAM 在 C 里到底怎么跑起来很多团队第一次把 SAM 往生产环境搬的时候都会卡在同一个地方Python 里segment_anything三行代码就能出掩码可一旦要嵌进 C 服务、要压延迟、要上边缘盒子整条链路就断了。PyTorch 权重是.pth推理框架是 libtorch部署侧却要求 TensorRT C中间还夹着 ONNX 导出、动态 shape、显存复用一堆事。这份资源要解决的就是这段断层——把 SAM 的编码器、提示编码器、掩码解码器拆开用 TensorRT 在 C 里重建推理管线让分割从「能跑」变成「能上线」。它适合两类人一类是已经会用 Python 调 SAM、现在要把能力落到 C 服务或端侧的工程师另一类是想搞懂 TensorRT 多输入动态 shape 模型怎么组织的部署同学。SAM 的结构天然适合拆成多引擎这也是后面所有优化的基础。下面从模型拆解讲到引擎构建再到 C 侧的内存与后处理最后把踩过的坑摊开。2. SAM 三段式结构与 TensorRT 引擎拆分策略SAM 不是一个单体网络它由 image encoder、prompt encoder、mask decoder 三部分组成三者的输入输出形态完全不同。如果按整图一次性导出成单个 ONNX动态轴会非常难处理而且每次换提示点都要重跑编码器延迟直接翻倍。所以第一步不是急着写 C而是想清楚怎么拆。2.1 为什么必须拆成三个引擎image encoder 是 ViT 主干输入固定为1x3x1024x1024输出是1x256x64x64的 image embedding。这部分计算量最大但只跟图像有关跟提示无关。prompt encoder 处理点、框、掩码提示输出稀疏和稠密两种 embedding。mask decoder 把 image embedding 和 prompt embedding 融合输出多个候选掩码和 IoU 分数。拆开的好处很直接同一张图只跑一次 encoder之后换任意多个提示点只重跑 prompt encoder 和 mask decoder。实测在 1024 分辨率下encoder 占整体耗时的八成以上拆开后交互式分割的响应能压到几十毫秒级。常见做法是把 encoder 单独导成一个引擎prompt encoder 和 mask decoder 合并成第二个引擎因为后两者都很轻合并能减少一次 kernel launch 和显存往返。提示SAM 官方权重有vit_b、vit_l、vit_h三个规格部署优先选vit_b精度损失在多数分割任务里可接受显存和延迟却友好得多。2.2 导出 ONNX 时的动态轴设置TensorRT 构建引擎前要先拿到 ONNX。导出脚本里最容易翻车的是动态轴没标对导致后面 profile 配不上。下面是我常用的导出片段针对vit_bimport torch from segment_anything import sam_model_registry sam sam_model_registry[vit_b](checkpointsam_vit_b_01ec64.pth) sam.eval() # 只导出 image encoder输入固定尺寸 dummy_img torch.randn(1, 3, 1024, 1024) torch.onnx.export( sam.image_encoder, dummy_img, sam_encoder.onnx, input_names[image], output_names[image_embedding], opset_version17, do_constant_foldingTrue, ) # prompt encoder mask decoder 合并导出提示点数量设为动态 dummy_embed torch.randn(1, 256, 64, 64) dummy_points torch.randn(1, 1, 2) # (batch, num_points, xy) dummy_labels torch.ones(1, 1, dtypetorch.int64) torch.onnx.export( sam.prompt_encoder, (dummy_points, dummy_labels), sam_prompt_encoder.onnx, input_names[points, labels], output_names[sparse_embedding, dense_embedding], opset_version17, dynamic_axes{ points: {1: num_points}, labels: {1: num_points}, }, )逻辑说明encoder 输入尺寸固定不设动态轴这样 TensorRT 能充分做层融合和 kernel 自动调优。prompt encoder 的num_points设为动态是因为交互式分割里用户点数不固定从 1 个点到几十个点都要支持。opset_version选 17 是因为 SAM 里用到了scatter、grid_sample等算子低版本 opset 导出会报不支持。参数说明dummy_points的形状是(batch, num_points, 2)最后一维是(x, y)归一化坐标dummy_labels是int641 表示前景点0 表示背景点。这两个输入的 dtype 必须和后面 C 侧绑定的一致否则 TensorRT 会在setInputShape时报类型不匹配。2.3 用 trtexec 快速验证 ONNX 可构建性在写 C 之前先用trtexec把 ONNX 跑一遍确认算子都被支持。这一步能省掉大量在 C 里调试构建失败的时间trtexec --onnxsam_encoder.onnx \ --saveEnginesam_encoder.engine \ --fp16 \ --workspace4096 \ --minShapesimage:1x3x1024x1024 \ --optShapesimage:1x3x1024x1024 \ --maxShapesimage:1x3x1024x1024逻辑说明--fp16开启半精度SAM 在 fp16 下精度损失很小速度提升明显。--workspace单位是 MBencoder 这种大模型建议给到 4096 以上否则某些层会因为显存不足回退到慢速实现。三个 shape 参数对固定输入来说写一样即可但必须写否则 trtexec 会按默认动态处理。参数说明如果构建时报某个算子不支持先看 TensorRT 版本。TensorRT 8.6 以上对grid_sample支持较好8.5 及以下可能需要插件。构建成功后 trtexec 会打印每层耗时重点看有没有异常慢的层那通常是回退到了 CPU 或非优化实现。3. C 侧 TensorRT 推理管线搭建与显存管理引擎文件拿到手接下来是在 C 里把它跑起来。这部分的核心不是调 API而是显存和内存的生命周期管理。TensorRT 的IExecutionContext不是线程安全的多线程服务里每个线程要独立 context但 engine 可以共享。3.1 引擎加载与 context 创建先看引擎加载和反序列化的骨架代码#include NvInfer.h #include fstream #include vector class TrtLogger : public nvinfer1::ILogger { void log(Severity severity, const char* msg) noexcept override { if (severity Severity::kWARNING) std::printf([TRT] %s\n, msg); } }; nvinfer1::ICudaEngine* loadEngine(const std::string path, TrtLogger logger) { std::ifstream file(path, std::ios::binary); if (!file.good()) return nullptr; 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); nvinfer1::IRuntime* runtime nvinfer1::createInferRuntime(logger); nvinfer1::ICudaEngine* engine runtime-deserializeCudaEngine(buffer.data(), size); runtime-destroy(); // engine 已独立于 runtime return engine; }逻辑说明反序列化后runtime就可以销毁engine 持有自己的权重和层信息不依赖 runtime。这一步很多人会漏掉runtime-destroy()导致每次加载引擎都泄漏一小块内存长时间运行的服务会慢慢涨上去。参数说明TrtLogger里只打印 warning 及以上级别info 级别在加载大引擎时会刷屏。如果排查构建问题临时把级别调到kVERBOSE但生产环境务必调回来。3.2 显存分配与输入输出绑定TensorRT 推理需要显存指针数组输入输出按 binding 顺序排列。下面这段是分配和绑定的核心void* buffers[4]; // encoder: 1 in 1 out; decoder: 2 in 2 out cudaMalloc(buffers[0], 1 * 3 * 1024 * 1024 * sizeof(float)); // image cudaMalloc(buffers[1], 1 * 256 * 64 * 64 * sizeof(float)); // embedding // 动态 shape 输入必须先设置维度 auto* inputTensor context-getTensorShape(points); context-setInputShape(points, nvinfer1::Dims{3, {1, numPoints, 2}});逻辑说明setInputShape必须在enqueueV3之前调用且每次点数变化都要重设。如果忘了设TensorRT 会用上一次的 shape导致结果错位但不报错这是最隐蔽的坑之一。输出 tensor 的显存大小依赖输入 shape动态输入下要在设完 shape 后用getTensorShape重新查询输出维度再分配。参数说明buffers数组的顺序要和 engine 里 binding 的索引一致可以用engine-getBindingIndex(points)查询。fp16 模式下显存按half类型算别用sizeof(float)否则会多分配一倍。3.3 掩码后处理与阈值选择mask decoder 输出的是低分辨率 logits需要上采样回原图尺寸再二值化。后处理里有两个参数直接影响效果mask_threshold和pred_iou_thresh。参数典型值作用调大后果mask_threshold0.0logits 二值化阈值掩码变小、漏分割pred_iou_thresh0.88候选掩码质量过滤候选变少、可能全滤掉stability_score_thresh0.95掩码稳定性过滤边缘区域被丢弃逻辑说明SAM 默认输出 3 个候选掩码按 IoU 分数排序。pred_iou_thresh用来过滤低质量候选但如果设太高小目标可能一个都留不下。我一般先设 0.88观察实际输出再微调。stability_score_thresh对边缘敏感做精细分割时可以降到 0.9。参数说明上采样用双线性插值即可SAM 官方也是这么做的。注意 logits 是 fp16 时上采样前先转 fp32否则插值精度不够边缘会出现锯齿。4. 避坑与排查SAM TensorRT 部署里最容易翻车的五件事这一章是我自己踩过、也帮别人排查过的真实问题按「现象 → 原因 → 解决」写。很多问题不报错只是结果不对排查起来很费时间。4.1 引擎构建成功但推理输出全零现象trtexec构建没报错C 里enqueueV3也返回成功但输出 embedding 全是 0。原因输入显存没有真正拷贝数据或者拷贝用了错误的 stream。常见于cudaMemcpyAsync传了默认 stream而推理在另一个 stream 上导致拷贝和计算没有同步。解决拷贝和推理用同一个 stream或者在拷贝后加cudaStreamSynchronize。更稳妥的做法是输入拷贝用cudaMemcpyAsync配合同一 stream推理结束后再同步一次。4.2 动态点数下输出 shape 不更新现象第一次用 1 个点推理正常第二次传 5 个点结果还是按 1 个点算的。原因只调了setInputShape但没有重新查询输出 tensor 的维度输出显存还是按旧 shape 分配的。解决每次setInputShape后用context-getTensorShape(outputName)拿到新维度重新cudaMalloc或复用足够大的显存块。复用的话要保证块大小不小于新维度所需。4.3 fp16 下掩码边缘出现明显锯齿现象fp32 引擎边缘平滑换 fp16 后掩码边缘呈阶梯状。原因mask decoder 输出的低分辨率 logits 在 fp16 下精度不足上采样时误差被放大。解决把 mask decoder 的输出层强制保持 fp32或者在 C 后处理里先把 fp16 转 fp32 再做双线性插值。TensorRT 里可以用setPrecision对特定层锁定精度。4.4 多线程下 context 竞争导致结果错乱现象单线程正常开多线程后偶发掩码错位或崩溃。原因IExecutionContext不是线程安全的多个线程共用一个 context 会互相覆盖输入 shape 和显存绑定。解决每个线程创建独立 contextengine 共享。context 创建开销很小不用省这一步。如果并发量很大用 context 池但每个 context 仍只能被一个线程同时使用。4.5 ONNX 导出时 opset 版本与 TensorRT 不匹配现象ONNX 能导出但 TensorRT 构建时报某算子不支持。原因opset 版本过高TensorRT 当前版本还没实现该版本的算子定义。解决SAM 导出建议 opset 17TensorRT 8.6 以上支持良好。如果必须用更高 opset先查 TensorRT 官方算子支持矩阵确认目标算子已实现。实在不行就降 opset或者对不支持的算子写插件。5. 进阶技巧用 CUDA Graph 把交互式分割延迟再压一截前面把管线跑通了但如果做的是交互式分割——用户点一下要立刻出掩码——那 prompt encoder mask decoder 这段的 kernel launch 开销就不能忽略。这两个引擎都很轻单次计算可能只有几毫秒但几十个 kernel 的 launch 开销加起来能占到一半。CUDA Graph 能把这一串 launch 录制成一个图之后一次提交开销直接摊薄。做法是在第一次推理时用cudaStreamBeginCapture和cudaStreamEndCapture把 prompt encoder 和 mask decoder 的 enqueue 过程录下来之后每次换提示点只要输入显存地址不变直接cudaGraphLaunch就行。注意动态 shape 下输入维度变化会导致 graph 失效所以点数变化时要重新录制或者把点数固定成几个档位分别录制。// 录制阶段固定点数后录制一次 cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal); contextPrompt-setInputShape(points, nvinfer1::Dims{3, {1, numPoints, 2}}); contextPrompt-enqueueV3(stream); contextDecoder-enqueueV3(stream); cudaStreamEndCapture(stream, graph); cudaGraphInstantiate(graphExec, graph, nullptr, nullptr, 0); // 推理阶段更新输入显存后直接 launch cudaMemcpyAsync(pointsDev, pointsHost, bytes, cudaMemcpyHostToDevice, stream); cudaGraphLaunch(graphExec, stream);逻辑说明录制时输入显存地址必须固定所以 points 的显存要预分配好每次只更新内容不换地址。如果点数档位不多比如 1、4、16 三个档就录三个 graph按实际点数选最近的档位多出来的点补零并在 labels 里标 -1 忽略。参数说明cudaStreamCaptureModeThreadLocal比默认模式更安全避免捕获到其他线程的操作。graph 实例化后可以反复 launch但输入内容要在 launch 前更新完且更新用的 stream 要和 launch 的 stream 一致否则会有竞态。验证方法很简单用trtexec或自己写个计时循环对比开 graph 前后的端到端延迟。我实测在vit_b上prompt decoder 这段从 12ms 降到 5ms 左右交互体验差别很明显。从那以后我每次做交互式分割部署都会先问一句点数档位能不能固定能固定就上 CUDA Graph不能固定就老老实实每次设 shape。希望帮到你。本文还有配套的精品资源点击获取
返回列表