企业官网建设流程全解析

简介：这份资源面向具备一定 C++ 与深度学习部署基础的开发者，聚焦于用 NVIDIA TensorRT 在 GPU 上高效推理 Segment Anything Model（SAM），解决原始 PyTorch 版 SAM 推理速度慢、显存占用高、难以落地到实际工程的问题。项目名为 SPEED-SAM-C++-TENSORRT，通过 TensorRT 引擎构建与 CUDA 优化，显著提升 GPU 利用率，适合图像分割、交互式抠图、自动标注等场景的工程化部署。压缩包共 20 个文件，约 71.22MB，包含 3 个 cpp 源文件与多个 h 头文件构成的核心推理代码，2 个 onnx 模型文件（SAM 编码器与掩码解码器），以及 jpg、png 示例图片和 txt、license 等辅助文件，并配有 CMakeLists.txt 便于编译。已有 989 人学习下载。读者可据此获得一套完整的 C++ TensorRT 推理工程，理解从 ONNX 模型到引擎序列化、CUDA 预处理与后处理的实现路径，并借助示例图片快速验证分割效果，为自研高性能视觉部署方案提供参考。

1. 用 C++ 和 TensorRT 把 SAM 跑成生产力：这条路到底值不值得走

如果你手上有一个 SAM（Segment Anything Model）的.pth权重，想把它塞进 C++ 服务里做实时分割，大概率会卡在同一个地方：Python 里三行代码就能出掩码，搬到 C++ 里却要面对 ONNX 导出、TensorRT 引擎构建、显存管理、预处理对齐这一整套链路。标题里的「使用 C++ 中的 TensorRT 实现 SAM」，说的就是把 Meta 这套分割模型从研究脚本变成可部署的推理服务这件事。它解决的不是「SAM 能不能用」，而是「SAM 能不能在 C++ 工程里稳定、低延迟地用」。适合两类人：一类是做视觉算法落地、需要把模型嵌进现有 C++ 框架的工程师；另一类是想搞懂 TensorRT 部署全流程、拿 SAM 当练手项目的进阶开发者。这条路有门槛，但走通之后，你会发现所有 Transformer 类视觉模型的部署套路都是相通的。

2. SAM 的 C++ TensorRT 部署链路：从权重到引擎

2.1 为什么不能直接把 .pth 丢给 TensorRT

TensorRT 不认识 PyTorch 的权重格式，它吃的是自己序列化出来的 plan 文件，或者 ONNX 这类中间表示。所以第一步永远是格式转换。SAM 的结构比普通 CNN 复杂，它有三个部分：图像编码器（Image Encoder，基于 ViT）、提示编码器（Prompt Encoder）、掩码解码器（Mask Decoder）。图像编码器是重计算的大头，一张 1024×1024 的图跑一次 ViT-H 大概要几百毫秒；提示编码器和掩码解码器很轻，但它们是动态输入的——点、框、掩码提示的维度不固定。

这就带来一个关键选型问题：要不要把整个 SAM 导成一个 ONNX。我的建议是拆开导。图像编码器单独导成一个引擎，输入固定为[1, 3, 1024, 1024]；提示编码器和掩码解码器合并导成另一个引擎，输入维度用动态 shape 处理。原因很直接：图像编码器对同一张图只需要跑一次，之后换不同的点提示可以复用它的输出 embedding。如果你把整个模型导成一个引擎，每次换提示都要重跑 ViT，延迟直接翻几倍。

常见做法是先用 Python 导出 ONNX，再用trtexec或 TensorRT 的 Python API 构建引擎。导出时注意 opset 版本，SAM 里的grid_sample、einsum这类算子对 opset 有要求，opset 17 是比较稳的选择。

2.2 导出 ONNX 的两个关键脚本

先看图像编码器的导出。下面这段代码在 Python 环境里跑，依赖torch、segment-anything和onnx：

import torch from segment_anything import sam_model_registry # 加载官方权重，vit_h 是精度最高但最重的版本 sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth") sam.eval() # 只取图像编码器，包一层固定输入尺寸 class ImageEncoder(torch.nn.Module): def __init__(self, sam): super().__init__() self.encoder = sam.image_encoder def forward(self, x): # 输出 shape: [1, 256, 64, 64] return self.encoder(x) encoder = ImageEncoder(sam).cuda() dummy = torch.randn(1, 3, 1024, 1024).cuda() torch.onnx.export( encoder, dummy, "sam_image_encoder.onnx", input_names=["image"], output_names=["embedding"], opset_version=17, do_constant_folding=True, )

这段代码的逻辑是：把image_encoder单独抽出来，用一个固定尺寸的 dummy 输入触发 tracing。opset_version=17是为了支持 ViT 里的 attention 算子。do_constant_folding=True会把能提前算的常量折叠掉，减小 ONNX 体积。导出后你会得到一个几百 MB 的 ONNX 文件，ViT-H 的编码器本身就大，这是正常的。

提示编码器和掩码解码器的导出稍微麻烦一点，因为要处理动态输入。核心思路是把点坐标、点标签、框坐标都作为输入传进去：

import torch from segment_anything import sam_model_registry from segment_anything.modeling import PromptEncoder, MaskDecoder sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth") sam.eval() class PromptAndMask(torch.nn.Module): def __init__(self, sam): super().__init__() self.prompt_encoder = sam.prompt_encoder self.mask_decoder = sam.mask_decoder def forward(self, image_embedding, point_coords, point_labels): # sparse embeddings 来自点和框，dense 这里先留空 sparse, dense = self.prompt_encoder( points=(point_coords, point_labels), boxes=None, masks=None, ) # 解码出掩码，输出低分辨率 logits low_res_masks, _ = self.mask_decoder( image_embeddings=image_embedding, image_pe=self.prompt_encoder.get_dense_pe(), sparse_prompt_embeddings=sparse, dense_prompt_embeddings=dense, multimask_output=False, ) return low_res_masks model = PromptAndMask(sam).cuda() emb = torch.randn(1, 256, 64, 64).cuda() coords = torch.randn(1, 1, 2).cuda() # 一个点提示 labels = torch.ones(1, 1, dtype=torch.int).cuda() torch.onnx.export( model, (emb, coords, labels), "sam_prompt_mask.onnx", input_names=["image_embedding", "point_coords", "point_labels"], output_names=["low_res_masks"], opset_version=17, dynamic_axes={ "point_coords": {1: "num_points"}, "point_labels": {1: "num_points"}, }, )

这里dynamic_axes把点数量标成动态维度，这样同一个引擎能处理 1 个点、5 个点甚至一个框（框用两个点表示）。multimask_output=False表示只输出一个掩码，如果你需要 SAM 默认的三个候选掩码，把它改成True，输出维度会多一维。注意prompt_encoder里的get_dense_pe()依赖图像尺寸，如果你的输入不是 1024，这里要同步改。

2.3 用 trtexec 构建引擎并验证

ONNX 有了，接下来构建 TensorRT 引擎。最省事的方式是trtexec，它随 TensorRT 一起安装：

# 图像编码器：固定 shape，开 FP16 trtexec --onnx=sam_image_encoder.onnx \ --saveEngine=sam_image_encoder_fp16.plan \ --fp16 \ --workspace=4096 # 提示+解码器：动态 shape，需要指定 optimization profile trtexec --onnx=sam_prompt_mask.onnx \ --saveEngine=sam_prompt_mask_fp16.plan \ --fp16 \ --minShapes=point_coords:1x1x2,point_labels:1x1 \ --optShapes=point_coords:1x4x2,point_labels:1x4 \ --maxShapes=point_coords:1x16x2,point_labels:1x16 \ --workspace=2048

--fp16让 TensorRT 在支持 FP16 的显卡上用半精度计算，速度通常能提升 30% 到 50%，精度损失对分割任务几乎看不出来。--workspace是构建时允许用的显存上限，单位 MB，ViT-H 建议给到 4096。动态引擎必须给minShapes、optShapes、maxShapes三个 profile，optShapes是你最常用的点数量，TensorRT 会针对它做优化。如果你的点数量经常超过 16，把maxShapes调大，但注意显存占用会跟着涨。

构建完成后，用trtexec --loadEngine=xxx.plan --shapes=...可以快速验证引擎能不能跑通、延迟大概多少。这一步别跳过，我见过太多人 ONNX 导出成功、引擎构建报错，最后发现是某个算子不支持 FP16。

3. C++ 侧推理代码：把引擎跑起来

3.1 TensorRT 运行时初始化与显存管理

C++ 里用 TensorRT 的核心流程是：反序列化 plan → 创建 execution context → 分配显存 → 绑定输入输出 → 执行。下面是一个最小可用的封装：

#include <NvInfer.h> #include <cuda_runtime_api.h> #include <fstream> #include <vector> #include <memory> class TrtEngine { public: TrtEngine(const std::string& planPath) { // 读取 plan 文件到内存 std::ifstream file(planPath, std::ios::binary); file.seekg(0, std::ios::end); size_t size = file.tellg(); file.seekg(0, std::ios::beg); std::vector<char> buffer(size); file.read(buffer.data(), size); // 创建 runtime 和 engine runtime_.reset(nvinfer1::createInferRuntime(logger_)); engine_.reset(runtime_->deserializeCudaEngine(buffer.data(), size)); context_.reset(engine_->createExecutionContext()); // 为每个输入输出分配显存 int n = engine_->getNbBindings(); buffers_.resize(n); for (int i = 0; i < n; ++i) { auto dims = engine_->getBindingDimensions(i); size_t bytes = volume(dims) * sizeof(float); cudaMalloc(&buffers_[i], bytes); bindingNames_[engine_->getBindingName(i)] = i; } } ~TrtEngine() { for (auto ptr : buffers_) cudaFree(ptr); } void* buffer(const std::string& name) { return buffers_[bindingNames_[name]]; } bool infer(cudaStream_t stream) { return context_->enqueueV2(buffers_.data(), stream, nullptr); } private: static size_t volume(const nvinfer1::Dims& d) { size_t v = 1; for (int i = 0; i < d.nbDims; ++i) v *= d.d[i]; return v; } nvinfer1::Logger logger_; std::unique_ptr<nvinfer1::IRuntime> runtime_; std::unique_ptr<nvinfer1::ICudaEngine> engine_; std::unique_ptr<nvinfer1::IExecutionContext> context_; std::vector<void*> buffers_; std::map<std::string, int> bindingNames_; };

这段代码做了三件事：反序列化引擎、创建执行上下文、按 binding 分配显存。volume函数把Dims转成元素个数，乘以sizeof(float)得到字节数。注意这里假设所有输入输出都是 FP32，如果你用了 FP16 引擎，输入输出 binding 可能还是 FP32（TensorRT 会自动转换），但内部计算是 FP16。enqueueV2是异步执行，配合 CUDA stream 可以做流水线。

显存管理是 C++ 部署里最容易翻车的地方。SAM 的 ViT-H 编码器输出 embedding 是[1, 256, 64, 64]，FP32 下就是 4MB，不大；但引擎内部的中间激活值可能占几百 MB。如果你同时跑多个实例，显存会爆。我的习惯是给每个引擎单独一个 context，不要共享，因为 context 之间切换有开销。

3.2 预处理与后处理：对齐 Python 的数值

C++ 侧最容易出错的不是推理本身，而是预处理。Python 里 SAM 的预处理是：resize 到 1024×1024（保持长宽比，短边补零）、归一化（减 mean 除 std）、转 CHW。C++ 里如果直接用 OpenCV 的resize，插值方式和 PyTorch 不一致，会导致 embedding 有细微偏差，最终掩码边缘对不上。

#include <opencv2/opencv.hpp> // 输入 BGR 图，输出归一化后的 NCHW float 数据 void preprocess(const cv::Mat& bgr, float* dst, int targetSize = 1024) { int h = bgr.rows, w = bgr.cols; float scale = static_cast<float>(targetSize) / std::max(h, w); int nh = static_cast<int>(h * scale); int nw = static_cast<int>(w * scale); cv::Mat resized; // 用 INTER_LINEAR 和 PyTorch 的 bilinear 对齐 cv::resize(bgr, resized, cv::Size(nw, nh), 0, 0, cv::INTER_LINEAR); // 放到 1024x1024 画布上，右下补零 cv::Mat canvas = cv::Mat::zeros(targetSize, targetSize, CV_8UC3); resized.copyTo(canvas(cv::Rect(0, 0, nw, nh))); // BGR -> RGB, 归一化, HWC -> CHW cv::Mat rgb; cv::cvtColor(canvas, rgb, cv::COLOR_BGR2RGB); rgb.convertTo(rgb, CV_32FC3, 1.0 / 255.0); std::vector<cv::Mat> channels(3); cv::split(rgb, channels); float mean[3] = {0.485f, 0.456f, 0.406f}; float stdv[3] = {0.229f, 0.224f, 0.225f}; for (int c = 0; c < 3; ++c) { channels[c] = (channels[c] - mean[c]) / stdv[c]; memcpy(dst + c * targetSize * targetSize, channels[c].data, targetSize * targetSize * sizeof(float)); } }

关键点是INTER_LINEAR和 PyTorch 的bilinear对齐，以及归一化参数必须和训练时一致。mean和std是 ImageNet 的标准值，SAM 用的就是这套。补零的位置也要注意，SAM 官方是右下补零，如果你补在左上，embedding 会偏。

后处理是把low_res_masks（通常是 256×256）上采样回原图尺寸，然后二值化。上采样同样用双线性，阈值一般取 0。如果你需要更精细的边缘，可以在原图分辨率上做一次条件随机场，但那是另一个话题了。

3.3 一个完整的单点提示推理流程

把上面的片段串起来，一次完整的「给一个点、出一张掩码」流程是这样的：

void segmentByPoint(const cv::Mat& image, float px, float py) { TrtEngine encoder("sam_image_encoder_fp16.plan"); TrtEngine decoder("sam_prompt_mask_fp16.plan"); // 1. 预处理图像 std::vector<float> input(3 * 1024 * 1024); preprocess(image, input.data()); // 2. 拷贝到显存，跑图像编码器 cudaMemcpy(encoder.buffer("image"), input.data(), 3 * 1024 * 1024 * sizeof(float), cudaMemcpyHostToDevice); encoder.infer(0); // 3. 准备点提示，坐标要归一化到 [0,1] float coords[2] = {px / image.cols, py / image.rows}; int labels[1] = {1}; // 1 表示前景点 cudaMemcpy(decoder.buffer("point_coords"), coords, 2 * sizeof(float), cudaMemcpyHostToDevice); cudaMemcpy(decoder.buffer("point_labels"), labels, sizeof(int), cudaMemcpyHostToDevice); // 4. 把编码器的输出 embedding 拷到解码器输入 cudaMemcpy(decoder.buffer("image_embedding"), encoder.buffer("embedding"), 256 * 64 * 64 * sizeof(float), cudaMemcpyDeviceToDevice); // 5. 跑解码器，拿回掩码 decoder.infer(0); std::vector<float> masks(256 * 256); cudaMemcpy(masks.data(), decoder.buffer("low_res_masks"), 256 * 256 * sizeof(float), cudaMemcpyDeviceToHost); // 6. 上采样回原图并二值化 cv::Mat maskLow(256, 256, CV_32F, masks.data()); cv::Mat maskFull; cv::resize(maskLow, maskFull, image.size(), 0, 0, cv::INTER_LINEAR); cv::Mat binary = maskFull > 0.0f; }

这段代码里，点坐标归一化到[0,1]是 SAM 的要求，不是像素坐标。labels里 1 是前景点，0 是背景点，如果你想排除某个区域，可以加一个 label 为 0 的点。cudaMemcpyDeviceToDevice那一步是把编码器输出直接拷到解码器输入，避免了绕回主机内存，这是性能优化的关键。实际工程里，编码器只需要对每张图跑一次，之后换点提示只跑解码器，所以你会把编码器的输出缓存起来。

4. 避坑与排查：SAM + TensorRT 部署的五个血泪教训

4.1 引擎构建报「unsupported operator」

现象：trtexec构建时报某个算子不支持，常见的是GridSample或Einsum。原因：TensorRT 版本和 ONNX opset 不匹配，或者该算子在你的 TensorRT 版本里没有 FP16 实现。解决：先确认 TensorRT 版本，8.5 以上对GridSample支持较好；如果 FP16 不支持，去掉--fp16用 FP32 构建，或者用--layerPrecisions单独指定某些层用 FP32。

4.2 掩码边缘和 Python 结果对不上

现象：C++ 出的掩码比 Python 小一圈，或者边缘锯齿明显。原因：预处理 resize 的插值方式不一致，或者补零位置不同。解决：把 C++ 预处理后的图存成 npy，和 Python 的预处理结果逐像素对比，差异应该小于 1e-3。重点检查INTER_LINEAR和补零的Rect起点。

4.3 动态 shape 引擎第一次推理特别慢

现象：动态引擎第一次跑某个点数量时延迟很高，第二次就正常了。原因：TensorRT 对动态 shape 需要在实际 shape 上做一次 tactic 选择，第一次是 warmup。解决：在服务启动时用optShapes对应的点数量跑几次 warmup，把 tactic 缓存住。别等到线上第一个请求来了才 warmup。

4.4 显存泄漏导致跑几十次后 OOM

现象：服务跑一段时间后cudaMalloc失败。原因：每次推理都新建 context 或忘记释放 buffer。解决：context 和 buffer 在初始化时创建一次，复用；如果必须动态创建，确保析构里cudaFree。用nvidia-smi观察显存曲线，正常应该是平的。

4.5 FP16 引擎精度下降导致小目标漏分割

现象：大目标分割正常，小目标掩码缺失。原因：FP16 的动态范围有限，小目标的 logits 值很小，被截断成 0。解决：对掩码解码器保持 FP32，只对图像编码器用 FP16。或者用--fp16构建后在 C++ 里对输出做一次阈值补偿，但更稳的做法是解码器用 FP32。

5. 进阶技巧：让 SAM 在 C++ 里跑得更快更稳

5.1 用 CUDA Graph 消除 kernel 启动开销

当你的点提示数量固定时，整个推理流程的 kernel 序列是确定的，这时候可以用 CUDA Graph 把一串 kernel 捕获成一个图，一次提交。对于解码器这种小模型，kernel 启动开销可能占延迟的一半。做法是在 warmup 阶段用cudaStreamBeginCapture和cudaStreamEndCapture把enqueueV2包起来，之后每次推理直接cudaGraphLaunch。注意 TensorRT 的enqueueV2在 capture 模式下要用enqueueV3或者确保没有同步操作。

5.2 批量处理多个点提示

如果你的场景是「一张图、多个点、要多个掩码」，别循环调用解码器。把点提示拼成 batch，比如 8 个点提示拼成[8, 1, 2]的 coords，一次推理出 8 个掩码。解码器很轻，batch 8 的延迟比 batch 1 高不了多少，但吞吐量翻 8 倍。前提是导出 ONNX 时把 batch 维度也标成动态。

5.3 引擎版本兼容性检查

TensorRT 的 plan 文件不是向前兼容的，8.x 构建的引擎在 10.x 上可能加载失败。如果你的部署环境 TensorRT 版本会变，要么每次重新构建引擎，要么在代码里做版本检查：

int32_t major, minor, patch; runtime_->getEngineVersion(&major, &minor, &patch); // 和构建时的版本对比，不一致就重新构建

我一般会在服务启动时打印引擎的版本信息，出问题时第一眼就能看到是不是版本不匹配。

5.4 一个我踩过的坑：别在析构里做同步

早期我写的封装在析构函数里调用了cudaStreamSynchronize，结果服务退出时偶尔卡死。原因是析构顺序不确定，stream 可能已经被销毁了。后来改成显式调用shutdown()方法，在里面做同步和释放，析构只做兜底。这个习惯帮我省了很多「玄学」崩溃。

部署 SAM 这类模型，最深的体会是：Python 里一行predictor.predict()背后，是 C++ 里几十个显存拷贝、shape 对齐和版本检查。但一旦跑通，你会发现换任何 ViT 类模型都是同一套骨架。我现在的习惯是，每接一个新模型，先花半天把 ONNX 导出和引擎构建的脚本固化下来，后面调 C++ 就是填空。希望帮到你。

本文还有配套的精品资源，点击获取