资讯动态

SAM图像分割模型TensorRT加速与C++推理部署实战

发布时间:2026/8/26 10:33:57 来源:尧图企业网站定制
简介深度学习模型从训练到上线推理性能始终是工程落地的核心挑战。尤其是像SAM这类基于Transformer的视觉分割模型结构复杂、计算量大在CPU和GPU上直接运行往往难以满足实时性要求。TensorRT作为NVIDIA推出的高性能推理优化引擎通过算子融合、显存复用和自动调优等手段能够在保持精度的前提下显著压缩模型延迟。而C推理框架则进一步降低了运行时的额外开销适合将模型嵌入到高并发的生产服务中。在实际应用中无论是智能抠图、遥感图像分割还是视频对象追踪都需要在分割质量与响应速度之间取得平衡。本文以SAM的部署优化为例完整介绍了从PyTorch模型导出ONNX、再转换为TensorRT engine到编写C推理服务的全过程并分享了坐标变换、动态输入、显存管理等多个工程踩坑经验为图像分割模型的工业级部署提供了一套可落地的参考方案。 去年接过一个图像分割服务的优化需求接口要求单张图控制在 300ms 内返回结果第一版直接跑出了 2 秒多的延迟。查下来瓶颈一点都不意外模型用的是 Segment Anything Model简称 SAM图像编码器在 CPU 上根本没法看上了 GPU 也快不到哪去。后来把整套推理链路换到 TensorRT 上用 C 重写才把端到端延迟压到几十毫秒量级。这篇就把这套方案从模型转换到 C 推理框架的完整思路写清楚代码框架可以直接拿去改。如果你正被 SAM 的部署性能困扰或者想把 SAM 接进现有的 C 服务里这篇文章应该能帮你省掉不少弯路。内容包含SAM 的瓶颈分析、PyTorch 转 ONNX 再转 TensorRT 的详细流程、C 推理框架的核心实现、性能调优方向以及我实际踩过的几个坑。1. 先搞清楚 SAM 为什么这么慢瓶颈不只在模型体积1.1 图像编码器占了 90% 以上的耗时SAM 的推理流程和常规分类检测模型不太一样它不是一个输入图像直接输出结果的单阶段模型而是拆成三块图像编码器Image Encoder把 1024x1024x3 的图像编码成 256x64x64 的特征图。这部分本质是一个 ViT如果你用的是 ViT-H 版本参数规模超过 600M前向计算量非常大。提示编码器Prompt Encoder把点、框、掩码这些提示条件编码成 token。这一块很轻计算量可以忽略不计。掩码解码器Mask Decoder基于图像特征和提示 token输出最终的分割掩码。它是一个轻量 Transformer耗时几个毫秒级别。也就是说SAM 的大头开销全在图像编码器上。在 PyTorch FP32 下ViT-B 编码器单张图在 V100 上大概要 500~800msViT-H 直接奔着 2 秒以上去了。掩码解码器再快也没用因为整个接口的时延上限被编码器锁死了。1.2 TensorRT 为什么能比 PyTorch 快这么多很多人对 TensorRT 的理解停留在它做了 INT8 量化但量化只是其中一个加速手段。TensorRT 真正强的地方在于图优化和 Kernel 自动调优。具体展开是三个层面第一算子融合。PyTorch 的算子粒度很细一个 Conv 后面接 BN 再接 ReLU在 GPU 上是三个独立 kernel 依次执行。每个 kernel 启动有开销中间结果还要写回显存再读出来显存带宽就这么被浪费了。TensorRT 会把 ConvBNReLU 融合成一个 kernelAttention 里的 QKV 矩阵乘也会合并处理Kernel 数量可能直接砍掉一半以上。第二显存复用。TensorRT 在做图优化时会分析整个计算图中每个张量的生命周期把生命周期不重叠的张量分配到同一块显存显存占用能压得很低。同时减少显存分配和释放的次数也就减少了 cudaMalloc 的耗时这类调用一次就要几十微秒在长链路里积少成多也很可观。第三静态 shape 优化。把输入 shape 固定下来之后TensorRT 可以在构建期就把每个 kernel 的具体尺寸定死从而对每个算子做更激进的 Kernel 自动调优。用动态 shape 时 kernel 要做很多运行时判断性能会打折。这么说吧用 PyTorch 做 SAM 推理相当于每次都在临时拼装积木而 TensorRT 是提前把积木拼好放在那跑的时候直接用。这也是为什么同样的 GPU、同样的 FP16 精度下TensorRT 可以比 PyTorch 快 3~5 倍的原因。2. 模型转换ONNX 导出与 TensorRT engine 构建全流程2.1 环境版本搭配这一组组合我踩过最稳TensorRT 部署最麻烦的就是环境版本匹配问题。版本不匹配的表现很诡异有时导出 ONNX 好好的转 engine 时报算子不支持有时 engine 构建成功跑推理时结果全错。我最终稳定在下面这套组合组件版本CUDA11.8cuDNN8.9TensorRT8.6.1PyTorch2.0.1ONNX1.14onnx-graphsurgeon0.3.27Ubuntu 系统版本20.04这套组合不是最新的但非常稳。TensorRT 8.6 对 ONNX 的算子覆盖已经很全了而且和 PyTorch 2.0 导出的 ONNX 配合得很好。如果你用更新的 TensorRT 10.x也没问题但注意 10.x 里很多 API 改名了网上搜到的老教程代码可能要微调。另外我强烈建议在 Docker 里做这套环境。TensorRT 的安装包和 CUDA 版本强绑定容器方式能保证整个团队复现环境一致。我用的是nvcr.io/nvidia/tensorrt:23.06-py3这个镜像里面自带 TensorRT 8.5.3稍微升一下级就行。2.2 图像编码器和掩码解码器要分开导出SAM 的官方仓库里带了scripts/export_onnx_model.py但这个脚本导出的是整条推理链路。实际部署时这样做有两个问题第一图像编码器是纯静态的输入只有一张图输出一个 embedding。而掩码解码器需要多个输入图像特征、点坐标、点标签、掩码等且这些输入支持动态变化。一个 ONNX 文件里同时处理静态和动态输入构建 engine 时很容易出问题。第二SAM 的实际使用模式是一张图 多个提示。图像只需要编码一次Feature Embedding 缓存在显存里后续每个提示直接拿去解码。分开导出才能实现这个模式把编码器的输出持久化解码器每次基于缓存推理。所以我会分别导出两个 ONNX 文件。图像编码器的导出代码很简单固定输入尺寸到 1024x1024import torch from segment_anything import sam_model_registry sam sam_model_registry[vit_b](checkpointsam_vit_b_01ec64.pth).cuda().eval() dummy_input torch.randn(1, 3, 1024, 1024).cuda() torch.onnx.export( sam.image_encoder, dummy_input, sam_image_encoder.onnx, opset_version17, input_names[input], output_names[image_embedding], do_constant_foldingTrue, dynamic_axesNone, # 完全静态 )如果导出时报scaled_dot_product_attention相关的算子不支持错误这是 PyTorch 2.x 用 SDPA 实现 Attention 导致的。解法是把 image_encoder 里的 Attention 改回手写版或者降级 PyTorch 到 1.13。我个人推荐手写 Attention 的方式代码改动很小class Attention(nn.Module): def forward(self, x): B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) q, k, v qkv.unbind(2) attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) x (attn v).transpose(1, 2).reshape(B, N, C) x self.proj(x) return x掩码解码器的导出相对复杂它接受多个输入。我按官方推理逻辑构造 dummy 输入import numpy as np import torch.nn.functional as F # 固定 batch1, 64x64 特征图, 3 个 mask 输出 dummy_embed torch.randn(1, 256, 64, 64).cuda() dummy_point_coords torch.randn(1, 1, 2).cuda() * 1024 dummy_point_labels torch.ones(1, 1).cuda().long() dummy_mask_input torch.zeros(1, 1, 256, 256).cuda() dummy_has_mask_input torch.zeros(1, 1).cuda() dummy_orig_im_size torch.tensor([[480, 640]]).cuda() torch.onnx.export( sam.mask_decoder, (dummy_embed, dummy_point_coords, dummy_point_labels, dummy_mask_input, dummy_has_mask_input, dummy_orig_im_size), sam_mask_decoder.onnx, opset_version17, input_names[image_embedding, point_coords, point_labels, mask_input, has_mask_input, orig_im_size], output_names[masks, iou_predictions, low_res_masks], dynamic_axes{ point_coords: {1: num_points}, point_labels: {1: num_points}, }, )这里有个细节point_coords的维度是[1, num_points, 2]point_labels是[1, num_points]。num_points 是动态的因为用户可能点 1 个点也可能点 5 个点。但如果你的业务场景里提示数量是固定的比如只支持单点提示那可以连这个动态轴都去掉性能还能再提升一点。mask_input是用于输入低分辨率掩码作为 prompt 的如果第一次推理没有任何掩码传全 1 或者全 0 都可以。SAMA 官方在无掩码时传的是torch.zeros。2.3 用 trtexec 构建 engine 的核心参数ONNX 导出成功之后先用onnxsim做一次简化能去掉一些冗余算子python -m onnxsim sam_image_encoder.onnx sam_image_encoder_sim.onnx然后就用 TensorRT 自带的trtexec命令构建 engine。图像编码器是静态输入参数很简单trtexec \ --onnxsam_image_encoder_sim.onnx \ --saveEnginesam_image_encoder.engine \ --fp16 \ --memPoolSizeworkspace:8G \ --verbose构建时间取决于 GPU 型号T4 上大概 3~5 分钟A100 上快一些。这里--memPoolSizeworkspace:8G控制的是构建期工作空间上限不是最终显存占用。TensorRT 8.5 以上版本用--memPoolSize老版本是--workspace注意区分。掩码解码器带动态坐标轴需要指定 profile 的范围trtexec \ --onnxsam_mask_decoder_sim.onnx \ --saveEnginesam_mask_decoder.engine \ --fp16 \ --minShapespoint_coords:1x1x2,point_labels:1x1 \ --optShapespoint_coords:1x1x2,point_labels:1x1 \ --maxShapespoint_coords:1x10x2,point_labels:1x10 \ --memPoolSizeworkspace:4G如果之前导出时把 num_points 固定了这里就可以不用三个Shapes参数。构建完成之后可以用trtexec --loadEngine跑一次离线测速确认 engine 能正常加载和推理trtexec --loadEnginesam_image_encoder.engine --fp16 --shapesinput:1x3x1024x1024这一步排错最有效。engine 加载失败基本就是 TensorRT 和 ONNX 算子不兼容推理结果明显异常则要怀疑精度选择或数据预处理。注意trtexec 自带的测速数据是纯 GPU 推理时间不包括图像解码、resize、归一化、CPUGPU 拷贝这些环节。真实接口延迟要等 C 工程写完再整体测。3. C 推理框架搭建持久化 embedding 才是 SAM 部署的灵魂3.1 Engine 封装与上下文管理C 侧对 TensorRT 的封装核心就是两类对象ICudaEngine负责描述模型结构和权重IExecutionContext负责持有推理过程中的运行时资源。一个很关键的点ICudaEngine可以被多个线程共享但IExecutionContext是线程不安全的一个 context 同时只能跑一个推理。如果你的服务要并发处理多个请求常见的做法是共享同一个 engine为每个工作线程单独创建 context或者用一个互斥锁串行化访问。我推荐的工程结构是这样的class TrtEngine { public: bool load(const std::string enginePath) { std::ifstream file(enginePath, std::ios::binary); std::vectorchar data(std::istreambuf_iteratorchar(file), {}); runtime_ std::unique_ptrnvinfer1::IRuntime(nvinfer1::createInferRuntime(logger_)); engine_ std::unique_ptrnvinfer1::ICudaEngine(runtime_-deserializeCudaEngine(data.data(), data.size())); if (!engine_) return false; return true; } std::unique_ptrnvinfer1::IExecutionContext createContext() { return std::unique_ptrnvinfer1::IExecutionContext(engine_-createExecutionContext()); } const nvinfer1::ICudaEngine* getEngine() const { return engine_.get(); } private: class Logger : public nvinfer1::ILogger { void log(Severity severity, const char* msg) noexcept override { if (severity Severity::kWARNING) std::cout [TRT] msg std::endl; } } logger_; std::unique_ptrnvinfer1::IRuntime runtime_; std::unique_ptrnvinfer1::ICudaEngine engine_; };deserializeCudaEngine是把之前构建好的 engine 文件直接加载进来省去了每次启动时重新做图优化的开销。一个 100MB 左右的 engine反序列化加载大概需要 1~2 秒这比重新构建快太多了生产环境务必用这种方式不要在服务启动时现构建。3.2 图像前处理resize、padding 与归一化SAM 的图像编码器输入是 1024x1024x3 的 RGB 图像但原始图像尺寸五花八门所以前处理要做三步等比缩放、padding 到 1024x1024、归一化。等比缩放这一条太重要了。如果直接把图像强制拉伸到 1024x1024画面中的物体比例会失真分割结果的边界会变形。正确做法是计算缩放比例scale min(1024 / w, 1024 / h)。按比例缩放长边变成 1024短边不足 1024 的部分用 0 padding。记录缩放比例和 padding 偏移量推理完成后做坐标回映射。我封装了一个简单的预处理函数struct PreProcessResult { cv::Mat tensorInput; // 1x3x1024x1024 的 float 张量 float scale; int padW, padH; }; PreProcessResult preprocess(const cv::Mat img) { const int targetSize 1024; int h img.rows, w img.cols; float scale std::min((float)targetSize / w, (float)targetSize / h); int newW (int)std::round(w * scale); int newH (int)std::round(h * scale); cv::Mat resized; cv::resize(img, resized, cv::Size(newW, newH), 0, 0, cv::INTER_LINEAR); int padW targetSize - newW; int padH targetSize - newH; cv::Mat padded; cv::copyMakeBorder(resized, padded, 0, padH, 0, padW, cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0)); cv::Mat rgb; cv::cvtColor(padded, rgb, cv::COLOR_BGR2RGB); cv::Mat floatImg; rgb.convertTo(floatImg, CV_32FC3); // SAM 的归一化参数 static const float mean[] {123.675f, 116.28f, 103.53f}; static const float std[] {58.395f, 57.12f, 57.375f}; cv::Mat normalized; std::vectorcv::Mat channels(3); cv::split(floatImg, channels); for (int i 0; i 3; i) { channels[i] (channels[i] - mean[i]) / std[i]; } cv::merge(channels, normalized); // HWC - CHW cv::Mat chw; cv::dnn::blobFromImage(normalized, chw); PreProcessResult result; result.tensorInput chw.clone(); result.scale scale; result.padW padW; result.padH padH; return result; }归一化参数mean和std是 SAM 训练时用的固定值直接用就行。这里用 OpenCV 做预处理优点是代码简洁、可读性好缺点是 CPU 操作占一部分耗时。如果 QPS 要求高可以把 resize 和归一化挪到 CUDA 上用自定义 kernel 做这块我后面在调优章节细说。3.3 提示编码与坐标变换提示编码是 SAM 里最容易出错的地方尤其是坐标变换。用户输入的提示点是在原始图像坐标系下的坐标但编码器内部使用的坐标系是 64x64 的特征图坐标系图像编码器输出 256 通道的 64x64 特征图中间隔着一个 1024 的输入坐标系和一次 padding。坐标变换公式是x_1024 (x_orig * scale) padW_off // 这里假设 padding 加在右侧和下方 x_64 x_1024 / 16其中scale是预处理时算出的缩放比例padW_off是 padding 导致的偏移。如果 padding 是加在右侧和下方即copyMakeBorder的左上角对齐方式那么坐标偏移就是 0只需要乘 scale 再除以 16。如果你的 padding 策略是居中填充那么偏移量要加上padW / 2和padH / 2。这块写错就会导致分割结果整体漂移而且表面上看不出逻辑错误排查起来非常难受。我在实际工程里还遇到过一个细节点坐标必须转成 float不能用 int。因为坐标经过scale缩放后可能是小数直接用 int 会把位置精度丢掉在细小物体上表现得尤其明显。推理前构建解码器的输入张量void buildDecoderInputs( std::vectorfloat pointCoords, std::vectorfloat pointLabels, const cv::Point2f point, bool isPositive, float scale) { // SAM 内部实现里会给坐标加 0.5 偏移 float x point.x * scale 0.5f; float y point.y * scale 0.5f; pointCoords {x, y}; pointLabels {isPositive ? 1.0f : 0.0f}; }这里加 0.5 是 SAM 官方实现的做法目的是让坐标落在像素中心。如果用深度学习框架里常见的grid_sample对齐方式偏移量可能不一样务必和训练时保持一致。3.4 掩码解码与后处理掩码解码器的输入包括图像 embedding、点坐标、点标签、掩码输入、orig_im_size。输出有三个掩码3x256x256、IOU 预测3 个值、低分辨率掩码3x256x256。SAMA 官方设计是输出 3 个候选掩码分别对应不同的细粒度层级。推理时选择 IOU 预测最高的那个掩码作为最终结果。不过如果你做过实验会发现IOU 第二高的掩码有时质量更好这是因为 SAMA 的掩码预测是按最可能排序设计的但在某些边缘场景下反而第二个更平滑。实际工程里可以两个都输出结合业务判断。后处理的第一步是取 IOU 最大的候选掩码然后做 Sigmoid再上采样回原始图像尺寸。我的实现如下cv::Mat postprocessMask( const float* masksOutput, // shape: 3x256x256 const float* iouOutput, // shape: 3 cv::Size origSize, float scale, int padW, int padH) { int bestIdx 0; for (int i 1; i 3; i) { if (iouOutput[i] iouOutput[bestIdx]) bestIdx i; } cv::Mat lowRes(256, 256, CV_32FC1, const_castfloat*(masksOutput bestIdx * 256 * 256)); cv::Mat sigmoided; cv::exp(-lowRes, sigmoided); // sigmoid 1 / (1 exp(-x)) sigmoided 1.0f / (1.0f sigmoided); // 先去掉 padding再上采样到原图尺寸 float newW origSize.width * scale; float newH origSize.height * scale; cv::Mat noPad sigmoided(cv::Rect(0, 0, (int)std::round(newW), (int)std::round(newH))).clone(); cv::Mat fullMask; cv::resize(noPad, fullMask, origSize, 0, 0, cv::INTER_LINEAR); return fullMask; }注意这里先去 padding 再 resize 到原始尺寸。如果直接对整个 256x256 的掩码 resizepadding 区域会混入黑色的 0 值导致掩码边缘出现暗色条纹。这个顺序在视觉上完全不一样我第一次实现时顺序反了分割结果边缘像蒙了一层黑边排查花了不少时间。3.5 持久化 Embedding 的缓存设计SAM 部署和其他模型最大的不同在于一张图片可以反复输入不同的提示点来获取不同区域的分割结果。图像编码器只跑一次得到 256x64x64 的 embedding这个 embedding 应该缓存在显存里后续每个提示直接送到掩码解码器。所以我设计了两个接口class SamInferencer { public: bool setImage(const cv::Mat img); cv::Mat segment(const cv::Point2f point, bool isPositive); private: // 编码器相关 TrtEngine encoderEngine_; // 解码器相关 TrtEngine decoderEngine_; // 缓存 float* deviceEmbedding_; bool embeddingValid_ false; float scale_ 0.f; int padW_ 0, padH_ 0; cv::Size origSize_; };setImage内部做前处理、跑编码器、把 embedding 留在显存里。segment接收提示坐标直接走解码器。这个设计下一次图像编码 N 次提示分割的总耗时是encoder_time N * decoder_time和 N 次独立调用 SAM 相比省掉了 N-1 次编码器耗时优势非常明显。还有一个细节如果图像切换了embedding 缓存要失效否则新的提示会拿去和旧图的特征做解码结果必然是一团乱。我在setImage开头直接覆盖缓存并置embeddingValid_ true同时在segment开头做一次有效性检查防止误用未初始化的显存。4. 性能对比与调优方向4.1 FP32/FP16/INT8 的实测差距部署完成后我做了一轮完整的性能对比硬件是单张 V100SAM 用的是 ViT-B 权重。这里给出的是我实测得到的大致范围不同 TensorRT 版本和驱动会有波动但量级可以参考。模型精度图像编码器耗时掩码解码器耗时PyTorchFP32680ms12msTensorRTFP32220ms5msTensorRTFP1688ms2msTensorRTINT852ms1.5ms从 FP32 到 FP16图像编码器加速了 2.5 倍而从 PyTorch FP32 到 TensorRT FP16 总共加速接近 8 倍。对分割任务来说FP16 的精度损失基本肉眼不可见掩码 IoU 下降不超过 1%这个代价完全可以接受。INT8 的话分割这类对边界敏感的任务风险要大一些尤其是 SAM 特征图直接决定掩码解码器的输入量化误差可能被解码器放大。如果需要 INT8建议用 TensorRT 的--calib校准模式准备一批有代表性的图像做校准不要直接跳过校准跑默认量化。如果业务对精度比较敏感FP16 是当前性价比最高的选择。4.2 显存池与异步流的调优TensorRT 推理的显存分配通常在 engine 构建阶段就已经规划好了但每次调用enqueueV2时仍可能涉及一些运行时中间缓冲区的分配。如果服务长时间运行频繁调用会产生显存碎片导致显存占用逐渐攀升。我在工程里做了两个优化第一复用输出缓冲区和临时张量。尽可能把编码器输出、解码器输入输出缓冲区在初始化时一次性分配好整个生命周期复用不反复cudaMalloc。第二使用cudaStream_t异步推理。enqueueV2本身是异步返回的CPU 提交完 kernel 后可以立即处理下一个请求的预处理GPU 和 CPU 的计算重叠起来。配合多线程吞吐能提升 30% 以上。cudaStream_t stream; cudaStreamCreateWithFlags(stream, cudaStreamNonBlocking); // 异步推理 context-enqueueV2(bindings, stream, nullptr); // 异步拷贝结果 cudaMemcpyAsync(hostOutputs, deviceOutputs, bytes, cudaMemcpyDeviceToHost, stream); // 在 stream 上等待结果也可以做更多异步操作 cudaStreamSynchronize(stream);4.3 batch 化的进一步思考上面的数据是 batch1 的场景。如果业务特征是多张图像同时请求、每个请求只做一次分割那就应该考虑把图像编码器 batch 化。TensorRT 的静态 batch 实现非常简单导出 ONNX 时输入维度改成[N, 3, 1024, 1024]每次推理时把 N 张图像拼成一个 batch。batch4 时单张图的平均耗时能比 batch1 再降低 30%~40%。代价是单次请求的延迟会上升因为要凑齐 batch 才能一起推理。如果采用动态 batch工程复杂度会提高不少要做请求队列按 batch 窗口聚合请求还要处理超时请求的单独推理。以我的经验先做静态 batch 就能满足大部分生产需求。动态 batch 是后续优化吞吐时才需要触碰的方向。5. 部署时最容易踩的坑5.1 坐标变换的偏差会让掩码整体漂移这是我在整个部署过程中花时间最多的问题。现象是提示点明明点在物体上输出的掩码却偏向一侧看起来像提示点和掩码对不上。排查下来是scale的计算方式前后不一致。预处理时我用的是min(target/w, target/h)但坐标变换时用的是另一个比例结果在非正方形图像上就出现了系统性偏差。坐标变换必须和预处理严格共用同一个scale和 padding 偏移参数不能各算各的。另外一点用户在前端点击的坐标如果是 CSS 像素要考虑和高清屏的缩放比。这块可以在接口层统一转换成原图坐标再进模型不要在模型层处理避免不同前端设备的差异影响分割准确性。5.2 TensorRT 输出 buffer 的数据类型陷阱TensorRT 的输出 buffer 类型取决于模型最后一层的输出类型。FP16 engine 不意味着所有输出都是 FP16有些层的输出会被 TensorRT 自动转换成 FP32。如果我在分配输出缓冲区时按 FP16 大小分配然后按 FP32 去解释结果就会是一堆乱码。稳妥的做法是序列化 engine 后遍历所有 binding用getBindingDataType查询每个 binding 的实际 dtype再据此分配缓冲区。for (int i 0; i engine-getNbBindings(); i) { auto dims engine-getBindingDimensions(i); auto dtype engine-getBindingDataType(i); size_t bytes 1; for (int d 0; d dims.nbDims; d) bytes * dims.d[d]; bytes * (dtype nvinfer1::DataType::kFLOAT) ? 4 : 2; // 分配对应大小的缓冲区 }5.3 动态坐标轴导致 engine 精度异常我用动态 num_points 轴时遇到过一个诡异问题第一次推理正常第二次换一个点数不同的提示结果就不对了。原因在于 TensorRT 动态 shape 需要每次推理前调用setInputShape指定当前输入的具体 shape。如果忘了设置TensorRT 会用 profile 里的optShapes来推导但实际的输入缓冲区大小不匹配就会读越界或者读到旧数据。正确做法是在每次推理前对所有动态输入调用一次context-setInputShape(point_coords, nvinfer1::Dims4{1, numPoints, 2});还有一个隐含问题setInputShape必须在使用enqueueV2之前调用且每次输入 shape 变化时都要重新调用。如果连续多次推理 shape 相同可以只在第一次设置省一点 CPU 开销。5.4 多线程并发时的 CUDA context 混乱接入 HTTP 服务后请求是并发的。如果多个线程同时调enqueueV2且共享同一个IExecutionContext会出现 CUDA context 的并发冲突轻则性能下降重则程序直接崩。解决思路有两种加全局互斥锁同一时刻只有一个线程能进入 TensorRT 推理。实现简单但会限制并发能力。每线程单独创建IExecutionContext共享同一个ICudaEngine。线程数控制在合理范围内N 个线程对应 N 个 context互不干扰。我推荐第二种。多线程推理工程上并不复杂难的是做好线程池和请求分发但这是通用架构问题和 TensorRT 本身无关。如果你只需要单路推理先用互斥锁方案起步也没问题等 QPS 瓶颈明确后再说。5.5 长期运行显存只增不减SAM 的 embedding 缓存在显存里是固定大小理论上不应该有显存增长。但实际跑长时间压力测试时发现显存还是缓慢上涨。排查发现是解码器里的动态坐标轴导致的。每次setInputShape后TensorRT 会根据当前 shape 分配临时缓冲区但如果下一次 shape 更大会分配更大的缓冲区而旧的缓冲区如果没有被显存复用机制回收就产生了内存碎片。解决办法是在构建 engine 时把maxShapes设置成实际能达到的上限比如maxNumPoints16这样 TensorRT 在构建期就会按最大可能分配显存运行时不再动态扩容。如果你的提示点数变化范围很大那就得在代码里定期销毁重建 context或者干脆固定最大点数。最后再分享两个小技巧第一个是关于解码器输出掩码的后处理技巧。很多人会直接把 256x256 的低分辨率掩码 resize 到原图尺寸但这样得到的掩码边缘锯齿严重。我后来在解码器和原始尺寸之间加了一步先 resize 到 1024x1024 的输入分辨率再做一次轻微的 GaussianBlur然后再 resize 到原图。这样边缘平滑很多视觉上是质的提升代价只是增加 1ms 左右的耗时。第二个是关于代码结构的建议。把图像编码器和掩码解码器拆成两个独立的模块接口设计成setImage segment模式这个决策让后续迭代轻松很多。后来团队要支持视频分割直接把视频帧逐帧调 setImage 就可以要支持批量分割就在 segment 里加一个循环。反过来如果你一开始就把它们耦合在一起后面每个新需求都要动核心逻辑越改越痛苦。TensorRT 版本迭代蛮快的上面这些代码在 8.6 上没有问题如果换了新版 TensorRT重点关注 API 变更和 ONNX 算子支持范围整体思路是通用的。本文还有配套的精品资源点击获取

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价