简介本资源是一套基于TensorRT加速的SAMSegment Anything Model大模型C部署完整实现面向具备C基础与深度学习推理经验的开发者解决视觉分割模型在生产环境高效落地的难题适用于边缘设备部署、工业质检、智能驾驶等实时性要求高的场景。压缩包共22个文件涵盖核心C源码.cpp/.h、模型导出与推理逻辑如export.h、sam_utils.h、跨平台构建配置CMakeLists.txt、Dockerfile.dev、中文部署指南README_zh_windows.md、Jupyter教程.ipynb、典型测试图像与动图truck.jpg、truck.gif及开发环境配置.vscode设置整体体积仅1.74MB轻量易集成。已有814人学习下载提供从ONNX模型转换、TensorRT引擎构建到C端完整推理链路的可运行代码与分步说明目录结构模块清晰含线程池ThreadPool.h、内存缓冲buffers.h等工程化设计显著降低大模型部署门槛。1. 为什么用 TensorRT 部署 SAM 大模型必须写 C 源码不是 Python 不行而是推理延迟、显存占用和工业级集成卡在这里SAMSegment Anything Model在图像分割任务中表现出色但原始 PyTorch 版本在实际部署中面临三个硬伤单图推理常超 300msGPU A10显存峰值突破 4.2GBFP16且无法直接嵌入 C 主控系统、边缘设备固件或工业视觉 SDK。TensorRT 是 NVIDIA 官方针对 CUDA 设备深度优化的推理引擎它能把 ONNX 导出的 SAM 模型编译成低延迟、低显存、高吞吐的序列化 engine 文件——但这一步本身不产生可调用接口真正落地必须手写 C 加载、预处理、推理、后处理全流程源码。这不是“为了炫技”而是因为Python 绑定如tensorrtpip 包缺乏对动态 batch、多输入张量image point prompt box prompt、以及 mask 解码逻辑的细粒度控制而工业场景要求毫秒级响应如 PCB 缺陷定位需 80ms、确定性内存占用避免 GC 波动导致产线停机、以及与 OpenCV/Qt/VisionPro 等 C 生态无缝对接。本文聚焦真实工程现场——从 ONNX 模型准备到最终segment()函数可被其他 C 模块直接调用每行代码都对应一个可验证的硬件行为。2. 构建 SAM 的 TensorRT 引擎ONNX 导出约束、动态 shape 设置与 profile 配置SAM 模型结构复杂其图像编码器ViT-H与提示编码器MLPTransformer存在多输入、多输出、动态分辨率依赖等特性。直接导出 ONNX 后若不做针对性处理TensorRT 编译会失败或生成非最优 engine。关键在于理解 SAM 的三类输入张量及其 shape 约束images:(1, 3, H, W)—— 图像尺寸必须为 64 倍数SAM 论文要求H/W ∈ [64, 1024]但 TensorRT 要求显式声明动态范围point_coords:(1, N, 2)—— 提示点坐标N 可变0~100需设为opt模式point_labels:(1, N)—— 对应标签0背景1前景-1无效同上动态维度。2.1 导出带 dynamic_axes 的 ONNX 模型Python 端import torch import onnx from segment_anything import sam_model_registry, SamPredictor # 加载官方 SAM 模型以 vit_h 为例 sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth) sam.eval() # 构造 dummy input —— 注意必须模拟最小/最大/最优三组尺寸 dummy_img torch.randn(1, 3, 1024, 1024) # 最大尺寸 dummy_points torch.tensor([[[0.5, 0.5], [0.7, 0.3]]], dtypetorch.float32) # 2 points dummy_labels torch.tensor([[1, 1]], dtypetorch.int32) # 导出 ONNX显式声明 dynamic_axes torch.onnx.export( sam.image_encoder, dummy_img, sam_image_encoder.onnx, opset_version17, do_constant_foldingTrue, input_names[images], output_names[image_embeddings], dynamic_axes{ images: {2: height, 3: width}, # H/W 动态 } ) # 导出 mask decoder含 prompt 输入 def forward_decoder(image_emb, points, labels): sparse_emb, dense_emb sam.prompt_encoder(pointspoints, labelslabels, boxesNone, masksNone) low_res_masks, iou_predictions sam.mask_decoder( image_embeddingsimage_emb, image_pesam.prompt_encoder.get_dense_pe(), sparse_prompt_embeddingssparse_emb, dense_prompt_embeddingsdense_emb, multimask_outputTrue ) return low_res_masks, iou_predictions # 使用 trace 方式导出 decoder因含 control flowscript 不稳定 traced_decoder torch.jit.trace( forward_decoder, (torch.randn(1, 256, 64, 64), dummy_points, dummy_labels), strictFalse ) torch.onnx.export( traced_decoder, (torch.randn(1, 256, 64, 64), dummy_points, dummy_labels), sam_mask_decoder.onnx, opset_version17, input_names[image_embeddings, point_coords, point_labels], output_names[low_res_masks, iou_predictions], dynamic_axes{ point_coords: {1: num_points}, point_labels: {1: num_points}, } )注意SAM 的prompt_encoder和mask_decoder必须分离导出。ViT 输出的image_embeddings尺寸为(1, 256, 64, 64)对应 1024×1024 输入该 shape 固定无需动态但 prompt 输入必须声明num_points动态轴否则 TensorRT 编译时无法生成支持任意点数的 engine。2.2 使用 trtexec 构建 engine 并验证 profile 范围TensorRT 7.2 要求显式定义 optimization profile尤其对height/width/num_points三类动态维度。以下命令生成支持64≤H,W≤1024且1≤N≤100的 engine# 编译 image encoder engine仅图像输入无 prompt trtexec \ --onnxsam_image_encoder.onnx \ --saveEnginesam_image_encoder.engine \ --minShapesimages:1x3x64x64 \ --optShapesimages:1x3x512x512 \ --maxShapesimages:1x3x1024x1024 \ --fp16 \ --workspace2048 \ --timingCacheFiletiming_cache.trt # 编译 mask decoder engine含 prompt 输入 trtexec \ --onnxsam_mask_decoder.onnx \ --saveEnginesam_mask_decoder.engine \ --minShapesimage_embeddings:1x256x64x64,point_coords:1x1x2,point_labels:1x1 \ --optShapesimage_embeddings:1x256x64x64,point_coords:1x10x2,point_labels:1x10 \ --maxShapesimage_embeddings:1x256x64x64,point_coords:1x100x2,point_labels:1x100 \ --fp16 \ --workspace1024 \ --timingCacheFiletiming_cache.trt参数含义实际取值建议说明--minShapes最小输入尺寸images:1x3x64x64,point_coords:1x1x2必须覆盖最小合法输入否则 runtime 报错--optShapes最常出现尺寸images:1x3x512x512,point_coords:1x10x2TensorRT 在此尺寸下做 kernel 选择与 memory layout 优化--maxShapes最大允许尺寸images:1x3x1024x1024,point_coords:1x100x2超出则 runtime 报错不可动态扩容提示--workspace2048单位是 MB设置过小会导致编译失败out of memoryA10 显存 24GB此处设 2048MB 是安全值。若使用 RTX 409024GB可增至--workspace4096以启用更多优化策略。3. C 核心推理类实现Engine 加载、内存管理与多输入同步绑定TensorRT C API 的核心难点不在推理本身而在生命周期管理engine、context、binding、device memory、stream 之间存在强依赖关系一处释放顺序错误即导致段错误。我们封装为SamInference类确保 RAII 安全。3.1 初始化与资源分配构造函数#include NvInfer.h #include NvInferRuntime.h #include opencv2/opencv.hpp #include vector #include memory class SamInference { private: std::unique_ptrnvinfer1::ICudaEngine m_engine; std::unique_ptrnvinfer1::IExecutionContext m_context; std::vectorvoid* m_bindings; // device pointers for inputs/outputs cudaStream_t m_stream; int m_input_idx, m_point_coord_idx, m_point_label_idx; int m_mask_out_idx, m_iou_out_idx; public: SamInference(const std::string encoder_engine_path, const std::string decoder_engine_path) { // Step 1: Load both engines from serialized files auto runtime nvinfer1::createInferRuntime(gLogger); std::ifstream enc_file(encoder_engine_path, std::ios::binary); std::ifstream dec_file(decoder_engine_path, std::ios::binary); std::vectorchar enc_buf((std::istreambuf_iteratorchar(enc_file)), {}); std::vectorchar dec_buf((std::istreambuf_iteratorchar(dec_file)), {}); m_engine std::unique_ptrnvinfer1::ICudaEngine( runtime-deserializeCudaEngine(enc_buf.data(), enc_buf.size()) ); m_context std::unique_ptrnvinfer1::IExecutionContext( m_engine-createExecutionContext() ); // Step 2: Allocate GPU memory for all bindings int nbBindings m_engine-getNbBindings(); m_bindings.resize(nbBindings); for (int i 0; i nbBindings; i) { auto dims m_engine-getBindingDimensions(i); size_t size 1; for (int d 0; d dims.nbDims; d) { size * dims.d[d]; } size_t elemSize sizeof(float); if (m_engine-getBindingDataType(i) nvinfer1::DataType::kINT32) { elemSize sizeof(int32_t); } cudaMalloc(m_bindings[i], size * elemSize); // Record binding index by name std::string name m_engine-getBindingName(i); if (name images) m_input_idx i; else if (name point_coords) m_point_coord_idx i; else if (name point_labels) m_point_label_idx i; else if (name low_res_masks) m_mask_out_idx i; else if (name iou_predictions) m_iou_out_idx i; } // Step 3: Create CUDA stream for async execution cudaStreamCreate(m_stream); } ~SamInference() { for (auto ptr : m_bindings) cudaFree(ptr); cudaStreamDestroy(m_stream); } };逻辑说明m_bindings存储的是 GPU 显存地址void*而非 host 内存。cudaMalloc分配大小时必须根据getBindingDimensions()获取每个 binding 的 shape并乘以sizeof(dtype)。SAM decoder 的point_labels是int32类型若误按float分配会导致越界写入。3.2 图像预处理与 device-to-device 数据拷贝SAM 要求输入图像归一化至[0,1]并减去 ImageNet mean/std且必须为 RGB 顺序。OpenCV 默认 BGR需转换cv::Mat preprocess_image(const cv::Mat img) { cv::Mat rgb; cv::cvtColor(img, rgb, cv::COLOR_BGR2RGB); // BGR → RGB cv::Mat float_img; rgb.convertScaleAbs(float_img, 1.0 / 255.0); // uint8 → float32 [0,1] // Apply normalization: (x - mean) / std const float mean[3] {0.485f, 0.456f, 0.406f}; const float std[3] {0.229f, 0.224f, 0.225f}; cv::Mat normed; float_img.convertScaleAbs(normed, 1.0f); for (int c 0; c 3; c) { cv::Mat channel normed(cv::Rect(c, 0, 1, normed.rows)); channel (channel - mean[c]) / std[c]; } return normed; } // Copy preprocessed image to GPU binding void copy_image_to_device(const cv::Mat img, void* device_ptr) { size_t img_size img.total() * img.elemSize(); cudaMemcpyAsync(device_ptr, img.data, img_size, cudaMemcpyHostToDevice, m_stream); }参数说明cudaMemcpyAsync必须与m_stream绑定否则多 batch 推理时可能因同步缺失导致结果错乱。img.total()返回像素总数img.elemSize()返回每个像素字节数CV_32F 为 4。3.3 执行 inference 并解析 mask 输出SAM decoder 输出low_res_masks形状为(1, 3, 256, 256)multimask_outputTrue需双线性上采样至原图尺寸并 sigmoidstd::vectorcv::Mat SamInference::infer( const cv::Mat image, const std::vectorstd::pairfloat, float points, const std::vectorint labels) { // 1. Preprocess image and copy to GPU cv::Mat normed preprocess_image(image); copy_image_to_device(normed, m_bindings[m_input_idx]); // 2. Prepare point inputs std::vectorfloat coords; std::vectorint32_t lables_int; for (size_t i 0; i points.size(); i) { coords.push_back(points[i].first); coords.push_back(points[i].second); lables_int.push_back(labels[i]); } cudaMemcpyAsync(m_bindings[m_point_coord_idx], coords.data(), coords.size() * sizeof(float), cudaMemcpyHostToDevice, m_stream); cudaMemcpyAsync(m_bindings[m_point_label_idx], lables_int.data(), lables_int.size() * sizeof(int32_t), cudaMemcpyHostToDevice, m_stream); // 3. Execute encoder → get image_embeddings m_context-setBindingShape(m_input_idx, nvinfer1::Dims4{1,3,normed.rows,normed.cols}); m_context-enqueueV2(m_bindings.data(), m_stream, nullptr); // 4. Execute decoder (requires image_embeddings from encoder) // Note: In practice, youd run encoder first, then feed its output to decoder // This example assumes encoder output is already bound to decoders first input // 5. Copy output back std::vectorfloat masks_host(1 * 3 * 256 * 256); cudaMemcpyAsync(masks_host.data(), m_bindings[m_mask_out_idx], masks_host.size() * sizeof(float), cudaMemcpyDeviceToHost, m_stream); cudaStreamSynchronize(m_stream); // 6. Convert to OpenCV Mat and upsample std::vectorcv::Mat masks; for (int i 0; i 3; i) { // 3 masks cv::Mat mask(256, 256, CV_32F, masks_host.data() i * 256 * 256); cv::Mat full_mask; cv::resize(mask, full_mask, image.size(), 0, 0, cv::INTER_LINEAR); cv::threshold(full_mask, full_mask, 0.0, 1.0, cv::THRESH_BINARY); masks.push_back(full_mask); } return masks; }关键点setBindingShape()必须在enqueueV2()前调用且 shape 必须与--optShapes中定义的范围一致。cudaStreamSynchronize(m_stream)是必须的同步点否则 host 端读取masks_host时数据尚未就绪。4. 部署实操VSCode CMake 构建环境配置与常见编译错误修复在 Windows 或 Linux 上构建 TensorRT C 项目最易卡在链接阶段。以下为 VSCode CMakeLists.txt 的最小可行配置适配 TensorRT 8.6.1CUDA 11.8。4.1 CMakeLists.txt 关键片段cmake_minimum_required(VERSION 3.10) project(SamTensorRT LANGUAGES CXX) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) # Find TensorRT find_package(TensorRT REQUIRED PATHS /usr/lib/x86_64-linux-gnu /opt/tensorrt/lib NO_DEFAULT_PATH) include_directories(${TENSORRT_INCLUDE_DIRS}) link_directories(${TENSORRT_LIBRARY_DIRS}) # Find OpenCV find_package(OpenCV REQUIRED) include_directories(${OpenCV_INCLUDE_DIRS}) # Executable add_executable(sam_inference main.cpp sam_inference.cpp) target_link_libraries(sam_inference ${TENSORRT_LIBRARIES} ${OpenCV_LIBS} cudart nvinfer nvinfer_plugin nvonnxparser nvparsers )注意find_package(TensorRT)在 Ubuntu 22.04 上常失败因 TensorRT 官方 deb 包未注册 cmake config。此时需手动指定路径find_package(TensorRT REQUIRED PATHS /usr/lib/x86_64-linux-gnu NO_DEFAULT_PATH)并确保/usr/lib/x86_64-linux-gnu/libnvinfer.so存在。4.2 VSCode tasks.json 配置Linux{ version: 2.0.0, tasks: [ { type: cppbuild, label: C/C: g build active file, command: /usr/bin/g, args: [ -g, ${file}, pkg-config --cflags opencv4, -I/opt/tensorrt/include, -L/opt/tensorrt/lib, -lnvinfer, -lnvonnxparser, -lnvparsers, -lcudart, -o, ${fileDirname}/${fileBasenameNoExtension} ], options: { cwd: ${fileDirname} }, problemMatcher: [$gcc], group: build, detail: Task generated by Debugger. } ] }4.3 三大高频编译/运行错误及修复方案错误现象根本原因修复命令/操作undefined reference to nvinfer1::IBuilder::createNetworkV2链接了旧版 TensorRT 库8.0sudo apt remove tensorrt sudo apt install tensorrt8.6.1.6-1cuda11.8UbuntuCUDA driver version is insufficient for CUDA runtime versionnvidia-smi显示驱动版本低于 CUDA 要求nvidia-smi查看驱动版本 → 下载匹配的 NVIDIA Driver →sudo ./NVIDIA-Linux-x86_64-*.run --no-opengl-filesSegmentation fault (core dumped)atcontext-executeV2()binding 数组未按 engine 的 binding order 排列在for (int i0; inbBindings; i)循环中不要假设 input/output 顺序必须用getBindingName(i)映射索引如if (nameimages) input_idxi;5. 性能调优与生产级验证batch 推理、显存复用与 mask 后处理加速单图推理满足不了产线节拍如 30FPS必须支持 batch 推理。但 SAM 的 prompt 输入天然不支持 batch不同图的点坐标数量不同因此采用dynamic batch padding策略将 N 张图拼成 batch对点坐标不足者补零labels 补-1SAM 自动忽略。5.1 支持 batch 的输入组织方式struct BatchInput { std::vectorcv::Mat images; // [N, H, W, 3] std::vectorstd::vectorstd::pairfloat,float all_points; std::vectorstd::vectorint all_labels; }; void prepare_batch_bindings(const BatchInput batch, SamInference infer) { // 1. Find max number of points size_t max_points 0; for (const auto pts : batch.all_points) max_points std::max(max_points, pts.size()); // 2. Flatten and pad all points/labels std::vectorfloat flat_coords(batch.images.size() * max_points * 2, 0.0f); std::vectorint32_t flat_labels(batch.images.size() * max_points, -1); for (size_t i 0; i batch.images.size(); i) { size_t offset i * max_points; for (size_t j 0; j batch.all_points[i].size(); j) { flat_coords[offset*2 j*2] batch.all_points[i][j].first; flat_coords[offset*2 j*2 1] batch.all_points[i][j].second; flat_labels[offset j] batch.all_labels[i][j]; } } // 3. Copy to GPU cudaMemcpyAsync(infer.m_bindings[infer.m_point_coord_idx], flat_coords.data(), flat_coords.size() * sizeof(float), cudaMemcpyHostToDevice, infer.m_stream); cudaMemcpyAsync(infer.m_bindings[infer.m_point_label_idx], flat_labels.data(), flat_labels.size() * sizeof(int32_t), cudaMemcpyHostToDevice, infer.m_stream); }参数说明max_points决定了 decoder engine 的num_points维度上限必须 ≤--maxShapes中设定值。padding 使用0.0和-1是 SAM 官方指定的无效值decoder 会自动 mask 掉。5.2 显存复用技巧避免重复 malloc/free在循环推理中频繁cudaMalloc/cudaFree会引入显著开销每次 ~0.1ms。解决方案是预分配一块大 buffer按需切片class GpuBufferPool { private: void* m_pool; size_t m_total_size; std::vectorstd::pairvoid*, size_t m_allocations; public: GpuBufferPool(size_t total_mb) : m_total_size(total_mb * 1024 * 1024) { cudaMalloc(m_pool, m_total_size); } void* allocate(size_t bytes) { // Simple first-fit allocator (for demo) for (auto alloc : m_allocations) { if (alloc.second 0) { alloc {m_pool, bytes}; return m_pool; } } // fallback: use pool head void* ptr static_castuint8_t*(m_pool) m_allocations.size() * 1024 * 1024; m_allocations.emplace_back(ptr, bytes); return ptr; } ~GpuBufferPool() { cudaFree(m_pool); } };提示实际项目中应使用cudaMallocAsynccudaMemPool_tCUDA 11.2但需 driver ≥ 465.19。对于 GTX 1070compute capability 6.1只能用传统cudaMalloc此时 buffer pool 是唯一高效方案。5.3 mask 后处理加速用 CUDA kernel 替代 OpenCV resizecv::resize在 CPU 上执行对 256×256→1024×1024 的双线性插值耗时约 1.2ms。改用 CUDA kernel 可压至 0.3ms__global__ void bilinear_upsample_kernel( const float* __restrict__ input, float* __restrict__ output, int in_h, int in_w, int out_h, int out_w) { int x blockIdx.x * blockDim.x threadIdx.x; int y blockIdx.y * blockDim.y threadIdx.y; if (x out_w || y out_h) return; float fx (float)x * in_w / out_w; float fy (float)y * in_h / out_h; int ix (int)floorf(fx); int iy (int)floorf(fy); float dx fx - ix; float dy fy - iy; float v00 input[iy * in_w ix]; float v10 (ix 1 in_w) ? input[iy * in_w ix 1] : v00; float v01 (iy 1 in_h) ? input[(iy 1) * in_w ix] : v00; float v11 (ix 1 in_w iy 1 in_h) ? input[(iy 1) * in_w ix 1] : v00; output[y * out_w x] v00 * (1 - dx) * (1 - dy) v10 * dx * (1 - dy) v01 * (1 - dx) * dy v11 * dx * dy; } // Launch dim3 block(16, 16); dim3 grid((out_w block.x - 1) / block.x, (out_h block.y - 1) / block.y); bilinear_upsample_kernelgrid, block(d_input, d_output, 256, 256, h, w);验证方法用nvprof --unified-memory-profiling off --metrics achieved_occupancy,flop_count_sp测量 kernel 占用率与计算强度确保 60% occupancy 且无 global load/store stall。本文还有配套的精品资源点击获取 TensorRT部署SAM大模型：C++推理工程实践与性能优化
简介本资源是一套基于TensorRT加速的SAMSegment Anything Model大模型C部署完整实现面向具备C基础与深度学习推理经验的开发者解决视觉分割模型在生产环境高效落地的难题适用于边缘设备部署、工业质检、智能驾驶等实时性要求高的场景。压缩包共22个文件涵盖核心C源码.cpp/.h、模型导出与推理逻辑如export.h、sam_utils.h、跨平台构建配置CMakeLists.txt、Dockerfile.dev、中文部署指南README_zh_windows.md、Jupyter教程.ipynb、典型测试图像与动图truck.jpg、truck.gif及开发环境配置.vscode设置整体体积仅1.74MB轻量易集成。已有814人学习下载提供从ONNX模型转换、TensorRT引擎构建到C端完整推理链路的可运行代码与分步说明目录结构模块清晰含线程池ThreadPool.h、内存缓冲buffers.h等工程化设计显著降低大模型部署门槛。1. 为什么用 TensorRT 部署 SAM 大模型必须写 C 源码不是 Python 不行而是推理延迟、显存占用和工业级集成卡在这里SAMSegment Anything Model在图像分割任务中表现出色但原始 PyTorch 版本在实际部署中面临三个硬伤单图推理常超 300msGPU A10显存峰值突破 4.2GBFP16且无法直接嵌入 C 主控系统、边缘设备固件或工业视觉 SDK。TensorRT 是 NVIDIA 官方针对 CUDA 设备深度优化的推理引擎它能把 ONNX 导出的 SAM 模型编译成低延迟、低显存、高吞吐的序列化 engine 文件——但这一步本身不产生可调用接口真正落地必须手写 C 加载、预处理、推理、后处理全流程源码。这不是“为了炫技”而是因为Python 绑定如tensorrtpip 包缺乏对动态 batch、多输入张量image point prompt box prompt、以及 mask 解码逻辑的细粒度控制而工业场景要求毫秒级响应如 PCB 缺陷定位需 80ms、确定性内存占用避免 GC 波动导致产线停机、以及与 OpenCV/Qt/VisionPro 等 C 生态无缝对接。本文聚焦真实工程现场——从 ONNX 模型准备到最终segment()函数可被其他 C 模块直接调用每行代码都对应一个可验证的硬件行为。2. 构建 SAM 的 TensorRT 引擎ONNX 导出约束、动态 shape 设置与 profile 配置SAM 模型结构复杂其图像编码器ViT-H与提示编码器MLPTransformer存在多输入、多输出、动态分辨率依赖等特性。直接导出 ONNX 后若不做针对性处理TensorRT 编译会失败或生成非最优 engine。关键在于理解 SAM 的三类输入张量及其 shape 约束images:(1, 3, H, W)—— 图像尺寸必须为 64 倍数SAM 论文要求H/W ∈ [64, 1024]但 TensorRT 要求显式声明动态范围point_coords:(1, N, 2)—— 提示点坐标N 可变0~100需设为opt模式point_labels:(1, N)—— 对应标签0背景1前景-1无效同上动态维度。2.1 导出带 dynamic_axes 的 ONNX 模型Python 端import torch import onnx from segment_anything import sam_model_registry, SamPredictor # 加载官方 SAM 模型以 vit_h 为例 sam sam_model_registry[vit_h](checkpointsam_vit_h_4b8939.pth) sam.eval() # 构造 dummy input —— 注意必须模拟最小/最大/最优三组尺寸 dummy_img torch.randn(1, 3, 1024, 1024) # 最大尺寸 dummy_points torch.tensor([[[0.5, 0.5], [0.7, 0.3]]], dtypetorch.float32) # 2 points dummy_labels torch.tensor([[1, 1]], dtypetorch.int32) # 导出 ONNX显式声明 dynamic_axes torch.onnx.export( sam.image_encoder, dummy_img, sam_image_encoder.onnx, opset_version17, do_constant_foldingTrue, input_names[images], output_names[image_embeddings], dynamic_axes{ images: {2: height, 3: width}, # H/W 动态 } ) # 导出 mask decoder含 prompt 输入 def forward_decoder(image_emb, points, labels): sparse_emb, dense_emb sam.prompt_encoder(pointspoints, labelslabels, boxesNone, masksNone) low_res_masks, iou_predictions sam.mask_decoder( image_embeddingsimage_emb, image_pesam.prompt_encoder.get_dense_pe(), sparse_prompt_embeddingssparse_emb, dense_prompt_embeddingsdense_emb, multimask_outputTrue ) return low_res_masks, iou_predictions # 使用 trace 方式导出 decoder因含 control flowscript 不稳定 traced_decoder torch.jit.trace( forward_decoder, (torch.randn(1, 256, 64, 64), dummy_points, dummy_labels), strictFalse ) torch.onnx.export( traced_decoder, (torch.randn(1, 256, 64, 64), dummy_points, dummy_labels), sam_mask_decoder.onnx, opset_version17, input_names[image_embeddings, point_coords, point_labels], output_names[low_res_masks, iou_predictions], dynamic_axes{ point_coords: {1: num_points}, point_labels: {1: num_points}, } )注意SAM 的prompt_encoder和mask_decoder必须分离导出。ViT 输出的image_embeddings尺寸为(1, 256, 64, 64)对应 1024×1024 输入该 shape 固定无需动态但 prompt 输入必须声明num_points动态轴否则 TensorRT 编译时无法生成支持任意点数的 engine。2.2 使用 trtexec 构建 engine 并验证 profile 范围TensorRT 7.2 要求显式定义 optimization profile尤其对height/width/num_points三类动态维度。以下命令生成支持64≤H,W≤1024且1≤N≤100的 engine# 编译 image encoder engine仅图像输入无 prompt trtexec \ --onnxsam_image_encoder.onnx \ --saveEnginesam_image_encoder.engine \ --minShapesimages:1x3x64x64 \ --optShapesimages:1x3x512x512 \ --maxShapesimages:1x3x1024x1024 \ --fp16 \ --workspace2048 \ --timingCacheFiletiming_cache.trt # 编译 mask decoder engine含 prompt 输入 trtexec \ --onnxsam_mask_decoder.onnx \ --saveEnginesam_mask_decoder.engine \ --minShapesimage_embeddings:1x256x64x64,point_coords:1x1x2,point_labels:1x1 \ --optShapesimage_embeddings:1x256x64x64,point_coords:1x10x2,point_labels:1x10 \ --maxShapesimage_embeddings:1x256x64x64,point_coords:1x100x2,point_labels:1x100 \ --fp16 \ --workspace1024 \ --timingCacheFiletiming_cache.trt参数含义实际取值建议说明--minShapes最小输入尺寸images:1x3x64x64,point_coords:1x1x2必须覆盖最小合法输入否则 runtime 报错--optShapes最常出现尺寸images:1x3x512x512,point_coords:1x10x2TensorRT 在此尺寸下做 kernel 选择与 memory layout 优化--maxShapes最大允许尺寸images:1x3x1024x1024,point_coords:1x100x2超出则 runtime 报错不可动态扩容提示--workspace2048单位是 MB设置过小会导致编译失败out of memoryA10 显存 24GB此处设 2048MB 是安全值。若使用 RTX 409024GB可增至--workspace4096以启用更多优化策略。3. C 核心推理类实现Engine 加载、内存管理与多输入同步绑定TensorRT C API 的核心难点不在推理本身而在生命周期管理engine、context、binding、device memory、stream 之间存在强依赖关系一处释放顺序错误即导致段错误。我们封装为SamInference类确保 RAII 安全。3.1 初始化与资源分配构造函数#include NvInfer.h #include NvInferRuntime.h #include opencv2/opencv.hpp #include vector #include memory class SamInference { private: std::unique_ptrnvinfer1::ICudaEngine m_engine; std::unique_ptrnvinfer1::IExecutionContext m_context; std::vectorvoid* m_bindings; // device pointers for inputs/outputs cudaStream_t m_stream; int m_input_idx, m_point_coord_idx, m_point_label_idx; int m_mask_out_idx, m_iou_out_idx; public: SamInference(const std::string encoder_engine_path, const std::string decoder_engine_path) { // Step 1: Load both engines from serialized files auto runtime nvinfer1::createInferRuntime(gLogger); std::ifstream enc_file(encoder_engine_path, std::ios::binary); std::ifstream dec_file(decoder_engine_path, std::ios::binary); std::vectorchar enc_buf((std::istreambuf_iteratorchar(enc_file)), {}); std::vectorchar dec_buf((std::istreambuf_iteratorchar(dec_file)), {}); m_engine std::unique_ptrnvinfer1::ICudaEngine( runtime-deserializeCudaEngine(enc_buf.data(), enc_buf.size()) ); m_context std::unique_ptrnvinfer1::IExecutionContext( m_engine-createExecutionContext() ); // Step 2: Allocate GPU memory for all bindings int nbBindings m_engine-getNbBindings(); m_bindings.resize(nbBindings); for (int i 0; i nbBindings; i) { auto dims m_engine-getBindingDimensions(i); size_t size 1; for (int d 0; d dims.nbDims; d) { size * dims.d[d]; } size_t elemSize sizeof(float); if (m_engine-getBindingDataType(i) nvinfer1::DataType::kINT32) { elemSize sizeof(int32_t); } cudaMalloc(m_bindings[i], size * elemSize); // Record binding index by name std::string name m_engine-getBindingName(i); if (name images) m_input_idx i; else if (name point_coords) m_point_coord_idx i; else if (name point_labels) m_point_label_idx i; else if (name low_res_masks) m_mask_out_idx i; else if (name iou_predictions) m_iou_out_idx i; } // Step 3: Create CUDA stream for async execution cudaStreamCreate(m_stream); } ~SamInference() { for (auto ptr : m_bindings) cudaFree(ptr); cudaStreamDestroy(m_stream); } };逻辑说明m_bindings存储的是 GPU 显存地址void*而非 host 内存。cudaMalloc分配大小时必须根据getBindingDimensions()获取每个 binding 的 shape并乘以sizeof(dtype)。SAM decoder 的point_labels是int32类型若误按float分配会导致越界写入。3.2 图像预处理与 device-to-device 数据拷贝SAM 要求输入图像归一化至[0,1]并减去 ImageNet mean/std且必须为 RGB 顺序。OpenCV 默认 BGR需转换cv::Mat preprocess_image(const cv::Mat img) { cv::Mat rgb; cv::cvtColor(img, rgb, cv::COLOR_BGR2RGB); // BGR → RGB cv::Mat float_img; rgb.convertScaleAbs(float_img, 1.0 / 255.0); // uint8 → float32 [0,1] // Apply normalization: (x - mean) / std const float mean[3] {0.485f, 0.456f, 0.406f}; const float std[3] {0.229f, 0.224f, 0.225f}; cv::Mat normed; float_img.convertScaleAbs(normed, 1.0f); for (int c 0; c 3; c) { cv::Mat channel normed(cv::Rect(c, 0, 1, normed.rows)); channel (channel - mean[c]) / std[c]; } return normed; } // Copy preprocessed image to GPU binding void copy_image_to_device(const cv::Mat img, void* device_ptr) { size_t img_size img.total() * img.elemSize(); cudaMemcpyAsync(device_ptr, img.data, img_size, cudaMemcpyHostToDevice, m_stream); }参数说明cudaMemcpyAsync必须与m_stream绑定否则多 batch 推理时可能因同步缺失导致结果错乱。img.total()返回像素总数img.elemSize()返回每个像素字节数CV_32F 为 4。3.3 执行 inference 并解析 mask 输出SAM decoder 输出low_res_masks形状为(1, 3, 256, 256)multimask_outputTrue需双线性上采样至原图尺寸并 sigmoidstd::vectorcv::Mat SamInference::infer( const cv::Mat image, const std::vectorstd::pairfloat, float points, const std::vectorint labels) { // 1. Preprocess image and copy to GPU cv::Mat normed preprocess_image(image); copy_image_to_device(normed, m_bindings[m_input_idx]); // 2. Prepare point inputs std::vectorfloat coords; std::vectorint32_t lables_int; for (size_t i 0; i points.size(); i) { coords.push_back(points[i].first); coords.push_back(points[i].second); lables_int.push_back(labels[i]); } cudaMemcpyAsync(m_bindings[m_point_coord_idx], coords.data(), coords.size() * sizeof(float), cudaMemcpyHostToDevice, m_stream); cudaMemcpyAsync(m_bindings[m_point_label_idx], lables_int.data(), lables_int.size() * sizeof(int32_t), cudaMemcpyHostToDevice, m_stream); // 3. Execute encoder → get image_embeddings m_context-setBindingShape(m_input_idx, nvinfer1::Dims4{1,3,normed.rows,normed.cols}); m_context-enqueueV2(m_bindings.data(), m_stream, nullptr); // 4. Execute decoder (requires image_embeddings from encoder) // Note: In practice, youd run encoder first, then feed its output to decoder // This example assumes encoder output is already bound to decoders first input // 5. Copy output back std::vectorfloat masks_host(1 * 3 * 256 * 256); cudaMemcpyAsync(masks_host.data(), m_bindings[m_mask_out_idx], masks_host.size() * sizeof(float), cudaMemcpyDeviceToHost, m_stream); cudaStreamSynchronize(m_stream); // 6. Convert to OpenCV Mat and upsample std::vectorcv::Mat masks; for (int i 0; i 3; i) { // 3 masks cv::Mat mask(256, 256, CV_32F, masks_host.data() i * 256 * 256); cv::Mat full_mask; cv::resize(mask, full_mask, image.size(), 0, 0, cv::INTER_LINEAR); cv::threshold(full_mask, full_mask, 0.0, 1.0, cv::THRESH_BINARY); masks.push_back(full_mask); } return masks; }关键点setBindingShape()必须在enqueueV2()前调用且 shape 必须与--optShapes中定义的范围一致。cudaStreamSynchronize(m_stream)是必须的同步点否则 host 端读取masks_host时数据尚未就绪。4. 部署实操VSCode CMake 构建环境配置与常见编译错误修复在 Windows 或 Linux 上构建 TensorRT C 项目最易卡在链接阶段。以下为 VSCode CMakeLists.txt 的最小可行配置适配 TensorRT 8.6.1CUDA 11.8。4.1 CMakeLists.txt 关键片段cmake_minimum_required(VERSION 3.10) project(SamTensorRT LANGUAGES CXX) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) # Find TensorRT find_package(TensorRT REQUIRED PATHS /usr/lib/x86_64-linux-gnu /opt/tensorrt/lib NO_DEFAULT_PATH) include_directories(${TENSORRT_INCLUDE_DIRS}) link_directories(${TENSORRT_LIBRARY_DIRS}) # Find OpenCV find_package(OpenCV REQUIRED) include_directories(${OpenCV_INCLUDE_DIRS}) # Executable add_executable(sam_inference main.cpp sam_inference.cpp) target_link_libraries(sam_inference ${TENSORRT_LIBRARIES} ${OpenCV_LIBS} cudart nvinfer nvinfer_plugin nvonnxparser nvparsers )注意find_package(TensorRT)在 Ubuntu 22.04 上常失败因 TensorRT 官方 deb 包未注册 cmake config。此时需手动指定路径find_package(TensorRT REQUIRED PATHS /usr/lib/x86_64-linux-gnu NO_DEFAULT_PATH)并确保/usr/lib/x86_64-linux-gnu/libnvinfer.so存在。4.2 VSCode tasks.json 配置Linux{ version: 2.0.0, tasks: [ { type: cppbuild, label: C/C: g build active file, command: /usr/bin/g, args: [ -g, ${file}, pkg-config --cflags opencv4, -I/opt/tensorrt/include, -L/opt/tensorrt/lib, -lnvinfer, -lnvonnxparser, -lnvparsers, -lcudart, -o, ${fileDirname}/${fileBasenameNoExtension} ], options: { cwd: ${fileDirname} }, problemMatcher: [$gcc], group: build, detail: Task generated by Debugger. } ] }4.3 三大高频编译/运行错误及修复方案错误现象根本原因修复命令/操作undefined reference to nvinfer1::IBuilder::createNetworkV2链接了旧版 TensorRT 库8.0sudo apt remove tensorrt sudo apt install tensorrt8.6.1.6-1cuda11.8UbuntuCUDA driver version is insufficient for CUDA runtime versionnvidia-smi显示驱动版本低于 CUDA 要求nvidia-smi查看驱动版本 → 下载匹配的 NVIDIA Driver →sudo ./NVIDIA-Linux-x86_64-*.run --no-opengl-filesSegmentation fault (core dumped)atcontext-executeV2()binding 数组未按 engine 的 binding order 排列在for (int i0; inbBindings; i)循环中不要假设 input/output 顺序必须用getBindingName(i)映射索引如if (nameimages) input_idxi;5. 性能调优与生产级验证batch 推理、显存复用与 mask 后处理加速单图推理满足不了产线节拍如 30FPS必须支持 batch 推理。但 SAM 的 prompt 输入天然不支持 batch不同图的点坐标数量不同因此采用dynamic batch padding策略将 N 张图拼成 batch对点坐标不足者补零labels 补-1SAM 自动忽略。5.1 支持 batch 的输入组织方式struct BatchInput { std::vectorcv::Mat images; // [N, H, W, 3] std::vectorstd::vectorstd::pairfloat,float all_points; std::vectorstd::vectorint all_labels; }; void prepare_batch_bindings(const BatchInput batch, SamInference infer) { // 1. Find max number of points size_t max_points 0; for (const auto pts : batch.all_points) max_points std::max(max_points, pts.size()); // 2. Flatten and pad all points/labels std::vectorfloat flat_coords(batch.images.size() * max_points * 2, 0.0f); std::vectorint32_t flat_labels(batch.images.size() * max_points, -1); for (size_t i 0; i batch.images.size(); i) { size_t offset i * max_points; for (size_t j 0; j batch.all_points[i].size(); j) { flat_coords[offset*2 j*2] batch.all_points[i][j].first; flat_coords[offset*2 j*2 1] batch.all_points[i][j].second; flat_labels[offset j] batch.all_labels[i][j]; } } // 3. Copy to GPU cudaMemcpyAsync(infer.m_bindings[infer.m_point_coord_idx], flat_coords.data(), flat_coords.size() * sizeof(float), cudaMemcpyHostToDevice, infer.m_stream); cudaMemcpyAsync(infer.m_bindings[infer.m_point_label_idx], flat_labels.data(), flat_labels.size() * sizeof(int32_t), cudaMemcpyHostToDevice, infer.m_stream); }参数说明max_points决定了 decoder engine 的num_points维度上限必须 ≤--maxShapes中设定值。padding 使用0.0和-1是 SAM 官方指定的无效值decoder 会自动 mask 掉。5.2 显存复用技巧避免重复 malloc/free在循环推理中频繁cudaMalloc/cudaFree会引入显著开销每次 ~0.1ms。解决方案是预分配一块大 buffer按需切片class GpuBufferPool { private: void* m_pool; size_t m_total_size; std::vectorstd::pairvoid*, size_t m_allocations; public: GpuBufferPool(size_t total_mb) : m_total_size(total_mb * 1024 * 1024) { cudaMalloc(m_pool, m_total_size); } void* allocate(size_t bytes) { // Simple first-fit allocator (for demo) for (auto alloc : m_allocations) { if (alloc.second 0) { alloc {m_pool, bytes}; return m_pool; } } // fallback: use pool head void* ptr static_castuint8_t*(m_pool) m_allocations.size() * 1024 * 1024; m_allocations.emplace_back(ptr, bytes); return ptr; } ~GpuBufferPool() { cudaFree(m_pool); } };提示实际项目中应使用cudaMallocAsynccudaMemPool_tCUDA 11.2但需 driver ≥ 465.19。对于 GTX 1070compute capability 6.1只能用传统cudaMalloc此时 buffer pool 是唯一高效方案。5.3 mask 后处理加速用 CUDA kernel 替代 OpenCV resizecv::resize在 CPU 上执行对 256×256→1024×1024 的双线性插值耗时约 1.2ms。改用 CUDA kernel 可压至 0.3ms__global__ void bilinear_upsample_kernel( const float* __restrict__ input, float* __restrict__ output, int in_h, int in_w, int out_h, int out_w) { int x blockIdx.x * blockDim.x threadIdx.x; int y blockIdx.y * blockDim.y threadIdx.y; if (x out_w || y out_h) return; float fx (float)x * in_w / out_w; float fy (float)y * in_h / out_h; int ix (int)floorf(fx); int iy (int)floorf(fy); float dx fx - ix; float dy fy - iy; float v00 input[iy * in_w ix]; float v10 (ix 1 in_w) ? input[iy * in_w ix 1] : v00; float v01 (iy 1 in_h) ? input[(iy 1) * in_w ix] : v00; float v11 (ix 1 in_w iy 1 in_h) ? input[(iy 1) * in_w ix 1] : v00; output[y * out_w x] v00 * (1 - dx) * (1 - dy) v10 * dx * (1 - dy) v01 * (1 - dx) * dy v11 * dx * dy; } // Launch dim3 block(16, 16); dim3 grid((out_w block.x - 1) / block.x, (out_h block.y - 1) / block.y); bilinear_upsample_kernelgrid, block(d_input, d_output, 256, 256, h, w);验证方法用nvprof --unified-memory-profiling off --metrics achieved_occupancy,flop_count_sp测量 kernel 占用率与计算强度确保 60% occupancy 且无 global load/store stall。本文还有配套的精品资源点击获取