ARTICLE DETAIL

资讯详情

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

TensorRT部署SAM实例:从PyTorch到C++推理的完整工程实践

TensorRT部署SAM实例:从PyTorch到C++推理的完整工程实践 简介用TensorRT部署SAM分割模型的C工程包面向需要将分割模型高效落地到NVIDIA GPU的开发者与研究者覆盖从模型转换、层融合、内核自动调优到推理执行的完整部署路径。压缩包共29个文件、约5.32MB包含C源码头文件与CPP、CMake构建配置、Jupyter教程、Markdown说明文档、JSON配置文件以及演示图片/动图等源码和文档分层存放便于按需查阅。已有245人学习下载属于偏向实战与流程拆解的入门进阶资料。其中不仅提供加载TensorRT优化模型的API调用、输入输出处理等关键环节示例还附带README、知识拓展文件及教程笔记可辅助理解SAM部署时的常见工程配置、模型导出与排错思路适合在真实C环境下集成推理模块时参考。1. 从“分割一切”到生产推理为什么必须用TensorRT拿到SAMSegment Anything Model第一反应是“分割一切”真香但真把它往生产环境里一放就傻眼了PyTorch 推理一帧 1024×1024 的图像在 T4 上要几百毫秒根本扛不住视频流。而 TensorRT 部署 SAM 的这套 C 代码与部署流程就是用来干掉这个性能瓶颈的。它的核心思路是把 SAM 的 image encoder 转成 TensorRT engine用 C 跑推理把整个分割管线压到几十毫秒级别。这套流程还覆盖了 mask decoder 的处理、动态 batch 和多精度模式适合所有想用 C 集成 SAM 的工程师——不管是做工业检测、图像编辑还是视频分割。下面我按从模型导出到推理验证的顺序把这套方案里的关键步骤和踩过的坑逐一写清楚。2. TensorRT部署SAM的完整流程从PyTorch权重到engine文件的四个关键步骤这一节我们解决“怎么把SAM变成TensorRT能用的engine”。整个转化链路是PyTorch权重 → ONNX → TensorRT engine。中间有四个步骤最容易出错我按顺序拆开讲。2.1 环境准备Ubuntu下TensorRT版本与CUDA/cuDNN的匹配先说环境。TensorRT不是一个孤立库它和CUDA、cuDNN、onnxparser强耦合。我一般用NGC的PyTorch容器作为基础镜像容器里已经配好匹配的CUDA和TensorRT版本。比如TensorRT 8.6对应CUDA 11.8TensorRT 9.0对应CUDA 12.0。如果版本不匹配最常见的现象是dynamic library加载失败或者在构建engine时报错Error: could not load library libnvinfer.so.8。我自己在Ubuntu 20.04上装过一次TensorRT 8.5用的是NVIDIA官方的deb包装了之后还需要额外装onnx-graphsurgeon因为SAM的encoder有一些自定义操作需要图表修改。这里给你一个参考版本表组件推荐版本说明CUDA11.8 / 12.0与TensorRT主版本匹配cuDNN8.6 / 8.9卷积加速版本过旧会导致INT8校准失败TensorRT8.6 / 9.08.4以下对SAM的ViT支持不完整ONNX1.14导出时用到较新的opset17操作PyTorch1.13SAM原版基于1.12建议1.13以上安装TensorRT本身不复杂主要坑在路径。deb包会把库装到/usr/lib/x86_64-linux-gnu而tar包需要手动把lib和include加到环境变量。我通常还会在~/.bashrc里加一行export LD_LIBRARY_PATH/opt/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH避免C编译时找不到动态库。2.2 导出ONNX把SAM的image encoder和mask decoder分别导出SAM是一个组合模型包含一个很重的image encoderViT-H以及一个轻量的mask decoder。部署时常见的做法是分开导出两个ONNXencoder负责把1024×1024的图像编码成embeddingdecoder负责根据prompt和embedding生成mask。为什么不导出成一个整体因为encoder只需要在图像变化时重新推理而decoder需要针对每个prompt快速推理。分开后可以缓存encoder的输出交互式点击时只跑decoder延迟能低一个数量级。这里给出导出image encoder的脚本# export_encoder.py import torch from segment_anything import sam_model_registry sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth) sam.eval() encoder sam.image_encoder # 直接取子模块跳过prompt路径 x torch.randn(1, 3, 1024, 1024) with torch.no_grad(): torch.onnx.export( encoder, x, sam_image_encoder.onnx, export_paramsTrue, opset_version17, input_names[images], output_names[image_embeddings], dynamic_axes{images: {0: batch}} # batch维度动态 )这段代码的关键在第7行sam.image_encoder直接拿到encoder模块绕过了SAM forward里对prompt的处理逻辑。导出的输入名是images输出名是image_embeddings这里建议固定大小1024×1024因为TensorRT后续如果用动态shapeheight和width变化会带来额外的一批优化候选其实收益不大反而增加构建时间。mask decoder的导出更麻烦一点因为它的输入有image_embeddings、point_coords、point_labels、mask_input。常见做法是把prompt数量固定为1个点然后把动态维度放在batch上这样C侧需要自己维护prompt数组。下面是一个可行的包装导出脚本# export_decoder.py import torch import torch.nn as nn from segment_anything import sam_model_registry class DecoderWrapper(nn.Module): def __init__(self, sam): super().__init__() self.sam sam def forward(self, image_embeddings, point_coords, point_labels, mask_input): sparse_embeddings, dense_embeddings self.sam.prompt_encoder( points(point_coords, point_labels), boxesNone, masksmask_input ) low_res_masks, iou_predictions self.sam.mask_decoder( image_embeddings, self.sam.prompt_encoder.get_dense_pe(), sparse_prompt_embeddingssparse_embeddings, dense_prompt_embeddingsdense_embeddings, ) return low_res_masks, iou_predictions sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth) sam.eval() wrapper DecoderWrapper(sam) embedding torch.randn(1, 256, 64, 64) point_coords torch.randn(1, 1, 2) point_labels torch.ones(1, 1, dtypetorch.long) mask_input torch.randn(1, 1, 256, 256) with torch.no_grad(): torch.onnx.export( wrapper, (embedding, point_coords, point_labels, mask_input), sam_mask_decoder.onnx, export_paramsTrue, opset_version17, input_names[image_embeddings, point_coords, point_labels, mask_input], output_names[masks, scores], dynamic_axes{ point_coords: {0: batch}, point_labels: {0: batch}, mask_input: {0: batch} } )这个包装类把prompt encoder和mask decoder合成一个模块输入就是encoder的输出和prompt输出低分辨率mask和iou分数。注意输出的masks是256×256的logits需要在C后面上采样到1024×1024并做sigmoid。导出后建议用onnx.checker验证一下然后马上用onnxruntime跑一次推理确保ONNX本身没问题再进TensorRT。这一步能省掉后面大量排查时间。我见过不少人跳过这步结果TensorRT报解析错误时不知道是模型问题还是转换问题。2.3 用trtexec构建engine精度模式与动态shape参数拿到ONNX后最快验证转换可行性的工具是trtexec它是TensorRT自带的命令行工具。对于image encoder我最常用的命令是trtexec --onnxsam_image_encoder.onnx \ --saveEnginesam_image_encoder_fp16.engine \ --fp16 \ --minShapesimages:1x3x1024x1024 \ --optShapesimages:1x3x1024x1024 \ --maxShapesimages:4x3x1024x1024 \ --workspace4096 \ --verbose参数说明--fp16开启FP16精度速度能提升一倍左右这也是工业部署的默认选择。--minShapes/--optShapes/--maxShapes定义动态shape的范围。这里我把batch维度设为1到4height和width固定防止TensorRT在推理时因为输入尺寸变化触发replanning。--workspace4096设置构建时的最大显存占用单位MB。SAM的ViT-H挺吃显存给小了会报Out of Memory。--verbose把构建日志打全方便确认每一个layer是否被正确解析。对于mask decoder由于输入数量多动态维度也复杂我建议用builder脚本而不是trtexec。原因很简单trtexec的命令行参数对多输入支持不好而且无法方便地设置多个优化profile。下面看自定义builder。2.4 自定义builder脚本批量转换与内存优化当模型输入多于两个或者要同时转换encoder和decoder时我一般写一个C/Python混合的builder。这里给一个纯C的builder核心代码展示如何设置动态shape和多profile// build_engine.cpp #include NvInfer.h #include NvOnnxParser.h #include fstream #include memory std::unique_ptrnvinfer1::ICudaEngine buildEngine( nvinfer1::IBuilder builder, nvinfer1::IBuilderConfig config, const char* onnxPath) { auto network builder.createNetworkV2(1); // 显式batch模式 nvonnxparser::IParser* parser nvonnxparser::createParser(*network, gLogger); parser-parseFromFile(onnxPath, static_castint(nvinfer1::ILogger::Severity::kINFO)); auto profile builder.createOptimizationProfile(); // 输入名必须和ONNX中的一致 profile-setDimensions(images, nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1,3,1024,1024)); profile-setDimensions(images, nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1,3,1024,1024)); profile-setDimensions(images, nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(4,3,1024,1024)); config.addOptimizationProfile(profile); config.setMemoryPoolSize(nvinfer1::MemoryPoolType::kWORKSPACE, 1ULL 32); // 4GB config.setFlag(nvinfer1::BuilderFlag::kFP16); return std::unique_ptrnvinfer1::ICudaEngine( builder.buildEngineWithConfig(*network, config)); }这段代码里有几个点值得说createNetworkV2(1)中的1表示显式batch模式这是动态shape的前提。setDimensions的三个参数分别对应最小、常规、最大profile。对于SAM这种输入固定在1024的模型最小和常规可以相等最大只放大batch。setMemoryPoolSize是新版TensorRT的API旧版叫setMaxWorkspaceSize。如果你用的是TensorRT 8.4以前的版本需要换回去。构建完成后可以把engine序列化到文件避免每次启动都重新构建std::ofstream file(sam_image_encoder_fp16.engine, std::ios::binary); file.write(engine-serialize().data(), engine-serialize().size());需要注意的是engine文件是和GPU架构绑定的。你在T4上构建的engine不能直接复制到A100上跑。这就是为什么很多部署流程里都在目标机器上现场构建。至此从权重到engine的链路就走通了。接下来进入C推理阶段这部分代码量最大也最容易在buffer管理上翻车。3. C推理实战加载engine并运行SAM的完整代码与参数说明构建好engine后剩下的事情全在C里。这一章我会给出一个能跑通的最小C工程CMake配置、engine加载、图像预处理、推理和后处理。3.1 项目结构与CMakeLists.txt配置我习惯把工程分成include/、src/、models/三个目录。models/放engine文件src/放主程序。下面是CMakeLists.txtcmake_minimum_required(VERSION 3.18) project(sam_trt) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(TENSORRT_ROOT /opt/TensorRT-8.6.1.6) include_directories(${TENSORRT_ROOT}/include) link_directories(${TENSORRT_ROOT}/lib) find_package(OpenCV REQUIRED) add_executable(sam_trt src/main.cpp) target_link_libraries(sam_trt nvinfer nvonnxparser cudart ${OpenCV_LIBS} )这里把TensorRT的路径写死成了/opt/TensorRT-8.6.1.6实际项目里建议用find_package或者环境变量动态获取。nvonnxparser只在构建engine时需要如果只跑推理可以去掉只留nvinfer和cudart。OpenCV用来做图像读取和resize预处理部分也可以用NPP代替但OpenCV在调试时更直观。3.2 加载engine反序列化与日志输出engine文件本质是一个序列化后的二进制对象运行时通过IRuntime反序列化。下面是加载代码// src/trt_engine.h #pragma once #include NvInfer.h #include fstream #include vector #include memory class InferLogger : public nvinfer1::ILogger { public: void log(Severity severity, const char* msg) override { if (severity ! Severity::kINFO) { std::cout [TensorRT] msg std::endl; } } }; inline std::unique_ptrnvinfer1::ICudaEngine loadEngine( const std::string enginePath, nvinfer1::ILogger logger) { std::ifstream file(enginePath, std::ios::binary); if (!file.good()) throw std::runtime_error(engine file not found); file.seekg(0, std::ios::end); size_t size file.tellg(); file.seekg(0, std::ios::beg); std::vectorchar data(size); file.read(data.data(), size); file.close(); auto runtime std::unique_ptrnvinfer1::IRuntime( nvinfer1::createInferRuntime(logger)); if (!runtime) throw std::runtime_error(create runtime failed); auto engine std::unique_ptrnvinfer1::ICudaEngine( runtime-deserializeCudaEngine(data.data(), size)); if (!engine) throw std::runtime_error(deserialize engine failed); return engine; }一个容易被忽略的坑是createInferRuntime返回的runtime必须保持存活直到engine被反序列化完毕。这个engine内部引用了一部分runtime的对象若在反序列化后立刻释放runtime某些版本下会导致未定义行为。我一般在整个生命周期都用unique_ptr持有runtime。3.3 图像预处理归一化与Letterbox适配分辨率SAM的输入要求是1024×1024而相机或测试图像不一定是这个尺寸。直接resize会破坏宽高比导致分割结果变形。所以预处理那里需要Letterbox把图像等比缩放然后填充到目标尺寸。填充值用0但要注意归一化方式。SAM原版预处理的归一化是像素值除以255用ImageNet的mean[0.485, 0.456, 0.406]和std[0.229, 0.224, 0.225]做标准化TensorRT输入要求的是CHW连续内存OpenCV给出的是HWC。这里我写一个完整预处理函数// src/preprocess.cpp #include opencv2/opencv.hpp #include vector void letterboxAndNorm(const cv::Mat src, float* dst, int targetSize) { int origW src.cols, origH src.rows; float scale std::min(static_castfloat(targetSize) / origW, static_castfloat(targetSize) / origH); int newW static_castint(origW * scale); int newH static_castint(origH * scale); cv::Mat resized; cv::resize(src, resized, cv::Size(newW, newH)); // 上下左右填充使最终尺寸为 targetSize x targetSize int padW (targetSize - newW) / 2; int padH (targetSize - newH) / 2; cv::Mat padded(targetSize, targetSize, src.type(), cv::Scalar(0, 0, 0)); resized.copyTo(padded(cv::Rect(padW, padH, newW, newH))); // 转Float并做RGB归一化 cv::Mat rgb; cv::cvtColor(padded, rgb, cv::COLOR_BGR2RGB); rgb.convertTo(rgb, CV_32FC3, 1.0 / 255.0); // HWC - CHW同时做标准化 const float mean[3] {0.485f, 0.456f, 0.406f}; const float stddev[3] {0.229f, 0.224f, 0.225f}; std::vectorcv::Mat channels(3); cv::split(rgb, channels); for (int c 0; c 3; c) { cv::Mat normalized; channels[c].convertTo(normalized, CV_32F, 1.0 / stddev[c], -mean[c] / stddev[c]); // 将每个通道复制到dst的CHW布局 memcpy(dst c * targetSize * targetSize, normalized.data, targetSize * targetSize * sizeof(float)); } }说明一下参数targetSize设为1024如果显存紧张可以用768或512但SAM的mask解码器在非1024分辨率下精度会下降不建议低于768。scale取两个比例的最小值这是Letterbox的标准做法保证图像完全落在画布内。padW和padH分别计算左右和上下的填充宽度这样即使非整数尺寸也能均匀分配。最后标准化时用了convertTo的alpha和beta参数一步完成(x - mean) / std减少循环开销。3.4 执行推理buffer绑定与后处理执行推理时需要把输入数据复制到CUDA显存然后调用executeV2或enqueueV2。SAM的encoder推理相对简单下面是核心代码// src/infer.cpp #include cuda_runtime_api.h #include NvInfer.h void runEncoder(std::unique_ptrnvinfer1::IExecutionContext context, float* inputHost, float* embeddingHost) { int inputIdx context-getEngine().getBindingIndex(images); int outputIdx context-getEngine().getBindingIndex(image_embeddings); nvinfer1::Dims outputDims context-getBindingDimensions(outputIdx); size_t outputSize 1; for (int i 0; i outputDims.nbDims; i) outputSize * outputDims.d[i]; void* inputDev; void* outputDev; cudaMalloc(inputDev, 1 * 3 * 1024 * 1024 * sizeof(float)); cudaMalloc(outputDev, outputSize * sizeof(float)); cudaMemcpy(inputDev, inputHost, 1 * 3 * 1024 * 1024 * sizeof(float), cudaMemcpyHostToDevice); // set tensor address新版TensorRT推荐方式 context-setTensorAddress(images, inputDev); context-setTensorAddress(image_embeddings, outputDev); bool success context-executeV2(0); // 0号stream单流推理 if (!success) throw std::runtime_error(execute failed); cudaMemcpy(embeddingHost, outputDev, outputSize * sizeof(float), cudaMemcpyDeviceToHost); cudaFree(inputDev); cudaFree(outputDev); }这里有几个细节容易踩坑setTensorAddress是TensorRT 8.5的API旧版本需要enqueueV2(bindings, stream, nullptr)传入一个void*数组。如果版本不对编译会报错。executeV2(0)中的0代表默认流如果生产环境有多个流要传入对应的cudaStream_t。embedding输出shape是[1, 256, 64, 64]前提是导出时固定了1024输入。若输入是动态batch输出shape也会变需要用getBindingDimensions去拿。decoder的推理与encoder原理一样但输入输出binding多不少。下面是一个精简的decoder推理片段输入是encoder得到的embedding、点坐标和标签void runDecoder(std::unique_ptrnvinfer1::IExecutionContext context, float* embedding, float* pointCoords, int* pointLabels, float* maskInput, float* maskOutput, float* scoreOutput) { void* devEmbedding, *devCoords, *devLabels, *devMaskIn, *devMaskOut, *devScore; cudaMalloc(devEmbedding, 1 * 256 * 64 * 64 * sizeof(float)); cudaMalloc(devCoords, 1 * 1 * 2 * sizeof(float)); cudaMalloc(devLabels, 1 * 1 * sizeof(int)); cudaMalloc(devMaskIn, 1 * 1 * 256 * 256 * sizeof(float)); cudaMalloc(devMaskOut, 1 * 1 * 256 * 256 * sizeof(float)); cudaMalloc(devScore, 1 * 1 * sizeof(float)); cudaMemcpy(devEmbedding, embedding, 1 * 256 * 64 * 64 * sizeof(float), cudaMemcpyHostToDevice); cudaMemcpy(devCoords, pointCoords, 1 * 1 * 2 * sizeof(float), cudaMemcpyHostToDevice); cudaMemcpy(devLabels, pointLabels, 1 * 1 * sizeof(int), cudaMemcpyHostToDevice); cudaMemcpy(devMaskIn, maskInput, 1 * 1 * 256 * 256 * sizeof(float), cudaMemcpyHostToDevice); context-setTensorAddress(image_embeddings, devEmbedding); context-setTensorAddress(point_coords, devCoords); context-setTensorAddress(point_labels, devLabels); context-setTensorAddress(mask_input, devMaskIn); context-setTensorAddress(masks, devMaskOut); context-setTensorAddress(scores, devScore); context-executeV2(0); cudaMemcpy(maskOutput, devMaskOut, 1 * 1 * 256 * 256 * sizeof(float), cudaMemcpyDeviceToHost); cudaMemcpy(scoreOutput, devScore, 1 * 1 * sizeof(float), cudaMemcpyDeviceToHost); // 释放所有dev指针 }注意maskInput初始化为全零它代表上一次预测的mask第一次调用时全零即可。pointLabels需要用int32或int64和导出时的dtype一致。如果pointLabels和pointCoords没有正确编码decoder输出会出现奇怪的mask这是我们下一章要讲的典型坑。后处理时maskOutput是256×256的logits需要上采样回原始输入尺寸再做sigmoid和阈值化。这里我用OpenCV做线性上采样// 后处理片段logits - binary mask cv::Mat logitsMat(256, 256, CV_32FC1, maskOutput); cv::Mat upsampled; cv::resize(logitsMat, upsampled, cv::Size(1024, 1024), cv::INTER_LINEAR); for (int i 0; i 1024 * 1024; i) { float val upsampled.atfloat(i / 1024, i % 1024); float sigmoid 1.0f / (1.0f std::exp(-val)); maskBinary[i] sigmoid 0.5f ? 255 : 0; }至此一个最小可运行的SAM-TensorRT推理链路已经完整了。但要做上线还需要在精度、延迟和吞吐之间调参下一章讲参数怎么调。4. 参数调优与性能平衡精度、动态shape和多流的取舍Engine构建参数的设置直接影响推理速度、显存占用和精度。这一章从实际调优角度给出我试过的参数组合和效果对比。4.1 精度模式FP32、FP16与INT8的差距TensorRT支持三种常见精度FP32、FP16、INT8。SAM这种大型ViT模型FP32意味着每个激活值占用4字节显存和带宽压力都很大。我用T4实测batch11024×1024输入精度延迟ms显存占用GB精度损失FP32822.8基准FP16311.6几乎不可见INT8181.1mask边界略粗糙FP16是我最推荐的部署精度尤其当encoder后面的decoder对数值不敏感时。INT8需要准备校准数据集SAM的校准数据一般选COCO里随机抽200张图。校准数据没选好可能出现mask整体偏灰或者某些目标丢失的现象。如果做医学图像分割建议谨慎用INT8因为边界细节要求高。开启FP16的副作用是precision一旦设错比如batch size在动态范围内变化某些layer会因精度溢出出现NaN。我当时就遇到过FP16在batch1正常batch4时输出全部是0。后来加了一个layer级别精度覆盖把关键层强制为FP32才解决。这里说明TensorRT的精度模式设置不是全局一刀切而是可以逐层覆盖的。可以用network-setPrecision对指定layer强制精度。4.2 动态shape的坑为什么我建议固定高宽SAM的encoder输入是正方形。如果动态shape允许高宽变化TensorRT会在运行时对卷积和attention层做replanning这个开销可能达到几毫秒甚至几十毫秒。我在AB测试中将输入从1024×1024改为512×512动态shape下的延迟反而比固定shape高一倍。原因就是TensorRT在每次shape变化时重新选择kernel这个选择过程比推理本身还贵。所以我把minShapes、optShapes、maxShapes都设置成同一个尺寸1024×1024只让batch维度动态。这样既能满足批处理需求又不会引入replanning。注意optShapes决定了性能优化目标它不一定是min或max的中点而是你最常用的shape。如果你大部分时候跑batch1就设batch1不要设成batch4否则TensorRT会优先优化batch4的性能batch1反而退化。4.3 多流推理用CUDA Stream提升吞吐量单帧31ms已经能跑30FPS但如果你想在一条GPU上同时处理多路视频就需要多流。TensorRT的IExecutionContext本身不是线程安全的习惯做法是每个线程创建独立的context共享同一个engine。每个context绑定不同的CUDA stream推理在同一块GPU上并发。一个典型的多流配置// 每路视频一个stream、一个context、一套buffer int numStreams 4; std::vectorcudaStream_t streams(numStreams); std::vectorstd::unique_ptrnvinfer1::IExecutionContext contexts(numStreams); for (int i 0; i numStreams; i) { cudaStreamCreate(streams[i]); contexts[i].reset(engine-createExecutionContext()); // 每个context设置对应的device buffer } // 每个工作线程循环 cudaMemcpyAsync(inputDev[i], inputHost[i], inputSize, cudaMemcpyHostToDevice, streams[i]); contexts[i]-enqueueV2(bindings[i].data(), streams[i], nullptr); cudaMemcpyAsync(outputHost[i], outputDev[i], outputSize, cudaMemcpyDeviceToHost, streams[i]);注意enqueueV2是异步的必须在同一个stream上同步后再读输出否则会读到上一帧的数据。很多新手在这里翻车把cudaMemcpy放到enqueue后而没有cudaStreamSynchronize导致输出一直是同一个数组。正确做法是在每个stream完成后调用cudaStreamSynchronize(streams[i])或者用cudaEvent做跨stream同步。我把4路1080p视频流同时跑FP16 encoder总体吞吐从单流30FPS提升到80FPS显存占用也才4.5GB说明多流确实划算。这里的经验是不要盲目创建几十个context显存会在setTensorAddress时暴露通常context数量等于GPU线程数或视频路数即可。5. 避坑与常见问题engine加载失败、mask全黑和显存爆炸的根因无论构建还是推理总有几个问题隔三差五出现。这里写几条我实际踩过、并且在这套部署流程里出现频率最高的坑。5.1 engine文件加载失败版本、序列化方式和路径问题现象运行时deserializeCudaEngine返回nullptr或者抛出Unsupport common type错误。原因最常见的是engine文件是由不同TensorRT版本构建的。TensorRT的engine格式并未完全向前兼容8.5的engine拿到8.6的runtime不一定能加载。其次序列化时只保存了engine内容没有保存设备和驱动信息跨GPU必定会失败。还有一种情况加载路径不对相对路径依赖当前工作目录而服务用systemd启动时工作目录是/找不到文件导致加载失败。解决第一构建和推理使用同一套TensorRT版本最好通过docker锁定。第二engine在目标机器现场构建不要跨机器拷贝。第三加载engine时用绝对路径并在启动时检查access(path, R_OK)提前暴露文件问题。5.2 输出mask全黑或尺寸完全不对预处理与后处理不匹配现象推理输出内存里全是0或者mask尺寸是奇怪的[1, 1, 256, 256]但不是预期的大小。原因大部分情况下是图像预处理尺寸和网络输入不匹配。SAM导出时固定了1024输入但你在C侧给的是512的tensorTensorRT执行时会直接在内部按stride读取越界内存轻则输出垃圾重则cuda错误。另一个常见原因是导出decoder时mask_input的shape和实际推理时使用的不一致。mask_input是上一轮预测的低分辨率mask初始化值应为全0形状与输出mask一致通常是256×256如果初始化成了[1,1,128,128]后处理输出的尺寸自然不对。解决固定输入尺寸C端强制resize到1024。对于decoder检查mask_input的创建逻辑用cudaMemset清零并确保shape与导出脚本一致。另外post-processing里的sigmoid是必须的吗如果导出时已经带上了sigmoid再执行一次就会导致所有输出变成0.5左右误以为是全黑。这两个操作只能二选一我的习惯是ONNX里不带sigmoidC里做logits到mask的转换这样灵活性更高。5.3 显存占用过高workspace与动态profile的叠加效应现象构建engine时提示Out of Memory或者推理时显存看起来是engine体积的好几倍。原因workspace参数是构建时的临时内存上限不代表engine最终占用的显存。实际运行时的显存占用是输入输出buffer engine运行需要的中间buffer。如果动态profile设置了多个shape每个shape都可能有一份优化后的中间buffer数量多起来显存会翻倍。另外没有显式释放CUDA buffer或者context创建太多都会导致显存爆炸。解决构建时setMemoryPoolSize设大一点4GB起步但运行时createExecutionContext不要超过GPU显存。我先用cudaMalloc一次性分配所有buffer然后复用避免每帧都malloc。再用nvidia-smi观察峰值反推需要的context数量。如果显存实在不够就降低batch上限或改用FP16。5.4enqueueV2卡死未同步stream或输入buffer未对齐现象程序卡在enqueueV2这一行不动CPU占用率也不变看起来像死锁。原因enqueueV2内部依赖当前stream是否空闲。如果你在同一个stream上同时排入了两个依赖关系不明的kernelTensorRT可能等待一个永远不会完成的先验kernel。更常见的是输入buffer地址没有对齐到512字节导致cudaMemcpy触发ub/channel错误。解决用cudaStreamSynchronize把上一个操作同步掉再调用下一轮推理。buffer采用cudaMalloc默认对齐即可不要用std::vector的data指针直接当CUDA指针。如果是在多线程环境还要注意不同线程共用一个stream那最好每个线程创建自己的stream。6. 验证方法用OpenCV可视化mask并接入视频流最后的落地验证是把推理结果画出来。最简单的方式是把mask输出转为单通道Mat再用imshow显示。下面这段代码把二值mask叠加到原图上// 后处理可视化 cv::Mat maskMat(1024, 1024, CV_8UC1, maskBinary); cv::Mat maskColor; cv::cvtColor(maskMat, maskColor, cv::COLOR_GRAY2BGR); maskColor.setTo(cv::Scalar(0, 0, 255), maskMat); // 红色mask cv::Mat original cv::imread(test.jpg); cv::resize(original, original, cv::Size(1024, 1024)); cv::addWeighted(original, 0.7, maskColor, 0.3, 0, original); cv::imwrite(result.jpg, original);这套可视化能快速判断预处理和后处理的参数对不对。我通常会准备一张包含人和背景的图片分别用点提示和框提示测试确认两套prompt路径都正常。接着再把整个推理封装成SamInfer类让外部只传入图片路径或cv::Mat返回mask。这样后续接入RTSP视频流时就只改数据来源不动推理逻辑。接入视频流的伪代码如下cv::VideoCapture cap(rtsp://your_stream); cv::Mat frame; while (cap.read(frame)) { std::vectorfloat input(3*1024*1024); letterboxAndNorm(frame, input.data(), 1024); // 运行encoder一次得到embedding runEncoder(contextEncoder, input.data(), embedding); // 固定点提示运行decoder得到mask runDecoder(contextDecoder, embedding, point, label, maskInput, maskOutput, scoreOutput); // 可视化并显示 }注意视频流场景下图像尺寸可能不是固定的letterbox参数需要根据每帧尺寸动态计算。如果分辨率变化频繁建议在letterboxAndNorm前先判断尺寸是否改变再决定是否重新分配buffer。否则频繁malloc会造成内存碎片。我在实际工程中还习惯将上面的推理函数用std::async包一层让它不阻塞主线程然后单独用一个渲染线程去显示或推流。从那以后我每次做完一个TensorRT模型都会强制跑一遍“固定输入 单流同步 可视化分支”的验证流程确保没有预处理和后处理的隐藏bug再进入并行优化。希望这些记录能帮你少走一些弯路。本文还有配套的精品资源点击获取
返回列表