PyTorch到TensorRT：SuperPoint+SuperGlue的C++部署实战

简介这份源码资源面向具备一定C与深度学习基础的计算机视觉开发者聚焦于用TensorRT加速SuperPoint特征点检测与SuperGlue特征匹配模型的推理部署可服务于SLAM、三维重建、相机姿态估计等对实时性要求较高的场景。压缩包共50个文件、约209.52MB包含5个cpp与10个h源文件、4个py转换脚本、3个engine与3个onnx模型文件以及png、gif等效果图与yaml、md配置说明覆盖从模型转换到C推理的完整链路。资源中附有运行环境说明涉及TensorRT、CUDA、cuDNN及OpenCV等依赖配置并给出室内外场景的engine与onnx权重便于对照验证。目前已有539人学习下载。读者可借此掌握层融合、精度校准等TensorRT优化思路理解SuperPoint与SuperGlue的工程化集成方式快速将模型从训练产物转化为可落地的推理引擎节省环境搭建与调试时间。1. 从 PyTorch 到 TensorRTSuperPointSuperGlue 的 C 部署到底在解决什么你手里有一份 SuperPoint SuperGlue 的 PyTorch 权重Python 下跑一张 640×480 的图特征提取加匹配大概 80150ms看着还行。但一旦要上嵌入式盒子、工控机或者 Jetson OrinPython 解释器、GIL、动态图调度这三座大山立刻把延迟推到 300ms 以上显存还飘。这个标题讲的就是把这条 Python 推理链路整体搬到 C TensorRT 上SuperPoint 出特征点和描述子SuperGlue 做图神经网络匹配中间用 TensorRT 引擎替换掉 PyTorch 的 forward最后用 C 串起来做成一个能塞进实际产品的推理程序。适合谁做 SLAM、视觉定位、AR 锚定、无人机视觉的工程师尤其是那些模型已经训好、Python 验证通过、现在卡在“怎么落地到 C 工程”这一步的人。热搜里“ubuntu 安装 tensorrt”“cpp 项目”“t4 1080p25帧”这些词恰好对应了这条链路上最常被搜的三个环节环境、工程结构、性能预期。2. 拆解 SuperPoint 与 SuperGlue 的 TensorRT 改造点2.1 两个模型各自的计算图特征SuperPoint 本质是一个 VGG 风格的编码器加两个解码头一个输出 65 通道的半稠密特征图64 个 cell 1 个 dustbin一个输出 256 维描述子。它的计算图非常规整全是 Conv、ReLU、MaxPool没有动态 shape没有控制流是 TensorRT 最喜欢的结构。SuperGlue 就不一样了它把关键点坐标和描述子拼成一个图经过多层 Attentional Graph Neural Network里面有大量的矩阵乘、softmax、以及基于关键点数量的动态索引。关键点数量 N 是变的这直接导致 TensorRT 的优化器在默认模式下会翻车——它需要固定 shape 才能做 kernel 融合和显存复用。所以改造的核心矛盾是SuperPoint 可以几乎无痛转 TensorRTSuperGlue 必须做“固定 N”或者“分档 N”的处理。常见做法是设定一个最大关键点数量 N_max比如 512 或 1024不足的补零多余的截断。补零会带来无效计算但换来的是引擎可以一次编译、反复使用。如果你追求极致可以按 N 分几档256/512/1024各编译一个引擎运行时按实际点数选档代价是显存占用成倍增加。2.2 为什么不用 ONNX Runtime 直接跑很多人第一反应是导出 ONNX然后用 ONNX Runtime 的 TensorRT Execution Provider。这条路能跑通但有两个硬伤。第一ONNX Runtime 的 TRT EP 在遇到 SuperGlue 里的动态索引操作时经常把整个子图回退到 CPU你看着用了 TensorRT实际一半算子在 CPU 上跑延迟反而比纯 PyTorch 还高。第二C 工程里引入 ONNX Runtime 会多一层依赖部署包体积大而且 TRT EP 的版本兼容性是个玄学——CUDA 版本、cuDNN 版本、TRT 版本、ONNX Runtime 版本四者必须严格对齐错一个就报 undefined symbol。直接写 TensorRT C API 虽然前期麻烦但引擎是自包含的部署时只需要 CUDA 运行时和 TensorRT 运行时干净得多。2.3 从 PyTorch 导出到 TensorRT 引擎的完整链路第一步把 PyTorch 模型导出成 ONNX。SuperPoint 导出时注意输入输出名字要固定SuperGlue 导出时需要把关键点数量维度设为动态轴。# export_superpoint.py import torch import torch.onnx from model.superpoint import SuperPoint # 假设你的模型定义在这里 # 加载权重 model SuperPoint(config{descriptor_dim: 256, nms_radius: 4}) model.load_state_dict(torch.load(superpoint.pth, map_locationcpu)) model.eval() # 构造 dummy input注意 SuperPoint 通常输入是灰度图shape 为 [B,1,H,W] dummy_input torch.randn(1, 1, 480, 640) torch.onnx.export( model, dummy_input, superpoint.onnx, input_names[image], output_names[scores, descriptors], opset_version11, dynamic_axes{ image: {0: batch, 2: height, 3: width}, scores: {0: batch, 2: height, 3: width}, descriptors: {0: batch, 2: height, 3: width} } )这段代码的关键在dynamic_axesSuperPoint 对输入尺寸是敏感的因为最后要 reshape 回特征图尺寸。如果你部署时输入分辨率固定比如就是 640×480可以把 dynamic_axes 去掉让 TensorRT 做全静态优化性能更好。opset_version11是保险选择再高一些的版本在 TensorRT 解析时可能遇到不支持的算子。SuperGlue 的导出更麻烦因为它的输入是keypoints和descriptors其中 keypoints 的 shape 是[B, N, 2]N 是动态的。# export_superglue.py import torch from model.superglue import SuperGlue model SuperGlue(config{ descriptor_dim: 256, weights: indoor, sinkhorn_iterations: 20, match_threshold: 0.2, }) model.eval() # N 设为动态轴但 TensorRT 最终需要固定这里先导出动态 ONNX dummy_kpts torch.randn(1, 512, 2) dummy_desc torch.randn(1, 256, 512) dummy_scores torch.randn(1, 512) torch.onnx.export( model, (dummy_kpts, dummy_desc, dummy_scores), superglue.onnx, input_names[keypoints, descriptors, scores], output_names[scores_out, matches], opset_version11, dynamic_axes{ keypoints: {1: num_keypoints}, descriptors: {2: num_keypoints}, scores: {1: num_keypoints}, scores_out: {1: num_keypoints}, matches: {1: num_keypoints} } )导出完成后用trtexec分别编译两个引擎。SuperPoint 可以直接用静态 shapetrtexec --onnxsuperpoint.onnx \ --saveEnginesuperpoint.engine \ --fp16 \ --workspace2048 \ --minShapesimage:1x1x480x640 \ --optShapesimage:1x1x480x640 \ --maxShapesimage:1x1x480x640SuperGlue 必须指定 N 的范围并且optShapes要设成你最常用的那个 Ntrtexec --onnxsuperglue.onnx \ --saveEnginesuperglue.engine \ --fp16 \ --workspace2048 \ --minShapeskeypoints:1x256x2,descriptors:1x256x256,scores:1x256 \ --optShapeskeypoints:1x512x2,descriptors:1x256x512,scores:1x512 \ --maxShapeskeypoints:1x1024x2,descriptors:1x256x1024,scores:1x1024注意--fp16不是万能的。SuperGlue 里的 softmax 和 sinkhorn 迭代对数值精度敏感FP16 可能导致匹配点数量骤降。我一般会先跑 FP32 引擎验证精度确认匹配结果和 PyTorch 一致后再试 FP16。如果 FP16 掉点严重就只对 SuperPoint 用 FP16SuperGlue 保持 FP32。3. C 工程结构从加载引擎到输出匹配点对3.1 最小可运行工程的目录与依赖一个能跑的 C 工程不需要多复杂但目录要清晰。我一般这样组织superpoint_superglue_trt/ ├── CMakeLists.txt ├── include/ │ ├── superpoint_trt.h │ ├── superglue_trt.h │ └── trt_utils.h ├── src/ │ ├── superpoint_trt.cpp │ ├── superglue_trt.cpp │ ├── trt_utils.cpp │ └── main.cpp ├── engines/ │ ├── superpoint.engine │ └── superglue.engine └── data/ └── test_pair/依赖只有三个CUDA Toolkit、TensorRT、OpenCV用来读图和可视化。CMakeLists.txt 里关键就几行cmake_minimum_required(VERSION 3.10) project(superpoint_superglue_trt) set(CMAKE_CXX_STANDARD 14) find_package(CUDA REQUIRED) find_package(OpenCV REQUIRED) include_directories( ${CUDA_INCLUDE_DIRS} ${OpenCV_INCLUDE_DIRS} /usr/local/TensorRT/include # 根据你的 TRT 安装路径调整 ${CMAKE_SOURCE_DIR}/include ) link_directories(/usr/local/TensorRT/lib) add_executable(match_demo src/main.cpp src/superpoint_trt.cpp src/superglue_trt.cpp src/trt_utils.cpp ) target_link_libraries(match_demo ${CUDA_LIBRARIES} ${OpenCV_LIBS} nvinfer nvinfer_plugin cudart )nvinfer和nvinfer_plugin是必须的后者在 SuperGlue 里如果用了自定义插件比如某些版本的 reshape 或 gather会用到。如果你在 Jetson 上部署TRT 路径通常是/usr/lib/aarch64-linux-gnu/头文件在/usr/include/aarch64-linux-gnu/需要相应调整。3.2 用 C 加载引擎并执行 SuperPoint 推理先写一个通用的 TRT 工具函数负责读引擎文件、反序列化、创建执行上下文。// trt_utils.cpp #include trt_utils.h #include fstream #include iostream std::vectorchar readEngineFile(const std::string path) { std::ifstream file(path, std::ios::binary); if (!file.good()) { std::cerr Engine file not found: path std::endl; return {}; } 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); file.close(); return buffer; } nvinfer1::ICudaEngine* loadEngine(const std::string path, nvinfer1::IRuntime* runtime) { auto buffer readEngineFile(path); if (buffer.empty()) return nullptr; auto engine runtime-deserializeCudaEngine(buffer.data(), buffer.size(), nullptr); if (!engine) { std::cerr Failed to deserialize engine: path std::endl; } return engine; }readEngineFile把整个引擎文件读进内存deserializeCudaEngine是 TRT 的标准反序列化接口。注意IRuntime要用createInferRuntime创建并且整个进程只创建一次多个引擎共享同一个 runtime。SuperPoint 的推理封装// superpoint_trt.cpp #include superpoint_trt.h #include cuda_runtime.h SuperPointTRT::SuperPointTRT(const std::string engine_path) { runtime_ nvinfer1::createInferRuntime(logger_); engine_ loadEngine(engine_path, runtime_); context_ engine_-createExecutionContext(); // 获取输入输出绑定索引 input_idx_ engine_-getBindingIndex(image); score_idx_ engine_-getBindingIndex(scores); desc_idx_ engine_-getBindingIndex(descriptors); // 分配显存尺寸根据引擎的 binding dimensions 确定 auto dims engine_-getBindingDimensions(input_idx_); input_size_ 1; for (int i 0; i dims.nbDims; i) input_size_ * dims.d[i]; cudaMalloc(input_buffer_, input_size_ * sizeof(float)); auto score_dims engine_-getBindingDimensions(score_idx_); score_size_ 1; for (int i 0; i score_dims.nbDims; i) score_size_ * score_dims.d[i]; cudaMalloc(score_buffer_, score_size_ * sizeof(float)); auto desc_dims engine_-getBindingDimensions(desc_idx_); desc_size_ 1; for (int i 0; i desc_dims.nbDims; i) desc_size_ * desc_dims.d[i]; cudaMalloc(desc_buffer_, desc_size_ * sizeof(float)); } void SuperPointTRT::infer(const float* input_cpu, float* score_cpu, float* desc_cpu) { cudaMemcpy(input_buffer_, input_cpu, input_size_ * sizeof(float), cudaMemcpyHostToDevice); context_-setTensorAddress(image, input_buffer_); context_-setTensorAddress(scores, score_buffer_); context_-setTensorAddress(descriptors, desc_buffer_); context_-enqueueV3(stream_); cudaStreamSynchronize(stream_); cudaMemcpy(score_cpu, score_buffer_, score_size_ * sizeof(float), cudaMemcpyDeviceToHost); cudaMemcpy(desc_cpu, desc_buffer_, desc_size_ * sizeof(float), cudaMemcpyDeviceToHost); }这里用的是 TRT 8.5 的setTensorAddressenqueueV3接口比老的enqueueV2更清晰。input_size_的计算假设输入是 NCHW 连续内存如果你的输入是灰度图记得在 CPU 侧把 OpenCV 的 BGR 转成 GRAY 并归一化到 [0,1]。3.3 SuperGlue 的输入构造与匹配结果解析SuperGlue 的输入不是原始图像而是 SuperPoint 输出的关键点坐标、描述子和置信度。所以 C 里需要先做一步后处理从 SuperPoint 的 score map 里提取关键点。// 从 score map 提取关键点简化版取 top-K void extractKeypoints(const float* score_map, int H, int W, std::vectorcv::Point2f keypoints, std::vectorfloat scores, int max_keypoints) { // score_map 是 [1, 65, H/8, W/8]需要先做 softmax 和 reshape // 这里省略 softmax 细节假设已经得到 [H/8, W/8] 的响应图 std::vectorstd::pairfloat, cv::Point2f candidates; for (int y 0; y H/8; y) { for (int x 0; x W/8; x) { float s score_map[y * (W/8) x]; if (s 0.01f) { // 阈值可调 candidates.emplace_back(s, cv::Point2f(x * 8 4, y * 8 4)); } } } std::sort(candidates.begin(), candidates.end(), [](auto a, auto b) { return a.first b.first; }); int n std::min((int)candidates.size(), max_keypoints); for (int i 0; i n; i) { keypoints.push_back(candidates[i].second); scores.push_back(candidates[i].first); } }这段代码是简化版实际 SuperPoint 的 score map 有 65 个通道需要先做 softmax 再取前 64 个通道的最大值并且要做 NMS。但核心逻辑就是把响应图变成一组带分数的关键点按分数排序取前 N 个。拿到关键点和描述子后构造 SuperGlue 的输入张量void prepareSuperGlueInput(const std::vectorcv::Point2f kpts, const std::vectorfloat desc, int desc_dim, int N_max, float* kpts_tensor, float* desc_tensor, float* scores_tensor) { int N std::min((int)kpts.size(), N_max); for (int i 0; i N; i) { kpts_tensor[i * 2] kpts[i].x; kpts_tensor[i * 2 1] kpts[i].y; for (int d 0; d desc_dim; d) { desc_tensor[d * N_max i] desc[i * desc_dim d]; } scores_tensor[i] 1.0f; // 或者用实际置信度 } // 补零到 N_max for (int i N; i N_max; i) { kpts_tensor[i * 2] 0; kpts_tensor[i * 2 1] 0; for (int d 0; d desc_dim; d) { desc_tensor[d * N_max i] 0; } scores_tensor[i] 0; } }注意描述子的内存布局SuperGlue 期望的是[1, desc_dim, N]也就是 desc_dim 在前N 在后。而 SuperPoint 输出的描述子通常是[1, desc_dim, H/8, W/8]你需要先按关键点位置把对应的描述子列取出来再转置成[desc_dim, N]。这一步在 C 里用循环做就行但要注意内存访问模式别搞成 cache 不友好。SuperGlue 的输出是scores_out和matches。matches是一个[1, N]的索引数组表示第一张图的第 i 个关键点匹配到第二张图的第matches[i]个关键点。如果matches[i] -1或者超出范围表示没有匹配。解析时直接遍历std::vectorcv::DMatch parseMatches(const float* matches, int N) { std::vectorcv::DMatch result; for (int i 0; i N; i) { int j (int)matches[i]; if (j 0 j N) { result.emplace_back(i, j, 0.0f); } } return result; }4. 避坑与排查部署 SuperPointSuperGlue 时最容易翻车的五个点4.1 引擎反序列化失败报 “serialization version mismatch”现象deserializeCudaEngine返回 nullptr日志显示版本不匹配。原因编译引擎的 TensorRT 版本和运行时的 TensorRT 版本不一致。比如你用 TRT 8.6 的trtexec编译但 C 链接的是 TRT 8.5 的libnvinfer.so。解决用dpkg -l | grep tensorrt确认版本确保编译和运行用的是同一套。如果必须在不同版本间迁移用trtexec --versionCompatible选项重新编译但会损失一些优化。4.2 SuperGlue 匹配点数量远少于 PyTorch现象同样的图对PyTorch 下匹配 200 对C 下只有 30 对。原因FP16 精度损失导致 sinkhorn 迭代不收敛或者输入描述子没有做 L2 归一化。SuperGlue 在计算相似度矩阵前会对描述子做内积如果描述子模长不是 1相似度尺度会偏。解决在 C 里对每个描述子做 L2 归一化并且先用 FP32 引擎验证。如果 FP32 正常、FP16 掉点就只对 SuperPoint 用 FP16。4.3 关键点坐标偏移匹配结果整体错位现象匹配点对在可视化时整体偏移几个像素。原因SuperPoint 的 score map 分辨率是输入图像的 1/8提取关键点时坐标要乘以 8 再加 0.5 的偏移因为 cell 中心在 4 左右。如果直接乘以 8 不加偏移会有半个 cell 的误差。解决x * 8 4或者x * 8 3.5具体看你的模型实现。这个偏移在 SLAM 里会导致位姿估计漂移血泪经验。4.4 显存泄漏跑几百帧后 OOM现象程序运行一段时间后cudaMalloc失败。原因每帧都创建新的IExecutionContext或者没有释放中间 buffer。解决IExecutionContext在初始化时创建一次反复使用所有cudaMalloc的 buffer 在析构函数里cudaFree。另外enqueueV3是异步的如果每帧都cudaStreamSynchronize不会泄漏但如果你用了多流要确保流之间同步正确。4.5 Jetson 上性能不升反降现象在 Jetson Orin 上TensorRT 版本比 PyTorch 还慢。原因Jetson 的 GPU 是统一内存架构cudaMemcpy的开销和桌面显卡不一样另外 Jetson 的 TRT 默认可能没有启用 DLA 或者没有做 INT8 校准。解决在 Jetson 上用trtexec --useDLACore0 --allowGPUFallback试试 DLA同时把输入分辨率降到模型实际需要的尺寸。如果还是慢检查是不是每帧都在做cudaMallocJetson 上这个操作特别贵。5. 进阶技巧用 CUDA Graph 把端到端延迟再压 30%当你把 SuperPoint 和 SuperGlue 都跑通之后会发现每帧的延迟里有一大块是 kernel launch 的开销。SuperPoint 有几十个 kernelSuperGlue 有上百个每个 kernel launch 大概 35 微秒加起来就是几百微秒。CUDA Graph 可以把这些 kernel 打包成一个图一次 launch 全部提交特别适合这种固定拓扑的推理链路。做法是先用一个 warmup 帧把 SuperPoint 和 SuperGlue 的enqueueV3都跑一遍然后用cudaStreamBeginCapture和cudaStreamEndCapture把整个序列捕获成一张图之后每帧只需要cudaGraphLaunch。// 捕获阶段 cudaStreamBeginCapture(stream_, cudaStreamCaptureModeGlobal); // SuperPoint 推理 context_sp_-setTensorAddress(image, input_buffer_sp_); context_sp_-setTensorAddress(scores, score_buffer_sp_); context_sp_-setTensorAddress(descriptors, desc_buffer_sp_); context_sp_-enqueueV3(stream_); // 中间的后处理关键点提取如果在 GPU 上做也可以捕获 // 这里假设后处理在 CPU 上所以图只包含两个引擎的推理 // SuperGlue 推理 context_sg_-setTensorAddress(keypoints, kpts_buffer_sg_); context_sg_-setTensorAddress(descriptors, desc_buffer_sg_); context_sg_-setTensorAddress(scores, scores_buffer_sg_); context_sg_-enqueueV3(stream_); cudaStreamEndCapture(stream_, graph_); // 之后每帧 cudaGraphLaunch(graph_exec_, stream_); cudaStreamSynchronize(stream_);注意CUDA Graph 捕获期间不能有cudaMemcpy到 host 的操作所以关键点提取如果放在 CPU 上就不能被捕获。解决办法是把关键点提取也写成 CUDA kernel或者接受图只覆盖两个引擎的推理后处理单独做。我实测在 T4 上640×480 输入、N512不用 Graph 时端到端 28ms用 Graph 后降到 19ms提升主要来自 kernel launch 的合并。另一个技巧是 SuperGlue 的 N 分档。如果你发现大部分帧的关键点数量在 300600 之间可以只编译一个 N512 的引擎超过 512 的截断不足的补零。截断会丢一些匹配但实际测试中前 512 个关键点已经覆盖了 95% 以上的有效匹配。如果你做的是高精度重建那就老老实实编译 N1024 的引擎显存换精度。最后说一个我自己的习惯每次改完模型或者换 TRT 版本一定先用trtexec跑一遍纯推理的 benchmark确认引擎本身的延迟和吞吐没问题再集成到 C 工程里。这样出问题时能快速定位是引擎的问题还是工程代码的问题。部署这件事后悔药就是提前把每一步的输入输出都 dump 出来对比别等到端到端跑不通了再回头查。希望帮到你。本文还有配套的精品资源点击获取