0. 简介

对于深度学习而言，通过模型加速来嵌入进C++是非常有意义的，因为本身训练出来的pt文件其实效率比较低下，在讲完BEVDET和FastBEV后，这里我们将展开实战，从pt到onnx再到tensorrt，以MixVpr作为例子，来向读者展示如何去跑CUDA版本的MixVpr，并给出相关的Tensorrt推理代码。这里最近受到优刻得的使用邀请，正好解决了我在大模型和自动驾驶行业对GPU的使用需求。UCloud云计算旗下的Compshare的GPU算力云平台。他们提供高性价比的4090 GPU，按时收费每卡1.88元，并附带200G的免费磁盘空间。暂时已经满足我的使用需求了，同时支持访问加速，独立IP等功能，能够更快的完成项目搭建。此外对于低性能的还有3080TI使用只需要0.88元，已经吊打市面上主流的云服务器了

对应的环境搭建已经在《如何使用共享GPU平台搭建LLAMA3环境（LLaMA-Factory）》、从BEVDET来学习如何生成trt以及如何去写这些C++内容介绍过了。对于自定义的无论是LibTorch还是CUDA这些都在《Ubuntu20.04安装LibTorch并完成高斯溅射环境搭建》这篇文章提到过了。这一章节我们来看一下怎么在平台上运行基于TensorRT的CUDA-FastBEV项目的。
在视觉定位任务中，重定位模块是至关重要的一环。尤其在深度学习中，重定位模块能够帮助系统精确地从庞大的数据库中快速找到与查询图像最匹配的图像。这个过程在自动驾驶、机器人导航以及增强现实等应用中尤为重要。通过使用深度学习模型，如 MixVPR，我们能够实现更加高效和精确的视觉重定位。

1. 安装环境

MIxVPR的环境安装还是比较简单的，只需要

pip intsall -r requirements.txt

MixVPR 使用的图像数据集可以从 GSV-Cities 仓库中获取。该数据集包含了来自多个城市的街景图像，适合作为视觉重定位任务的训练和测试数据集。

在本次实验中，我们选用了官方提供的预训练模型 resnet50_MixVPR_128_channels(64)_rows(2).ckpt。这是基于 ResNet50 架构的模型，并结合了 MixVPR 的特征聚合模块，能够高效地提取图像特征。

2. 导出 ONNX 模型

导出 ONNX 模型是实现跨平台推理的关键步骤。通过导出模型为 ONNX 格式，我们可以将模型应用于不同的硬件加速平台，如 TensorRT。

我们这里选用的是官方训练的resnet50_MixVPR_128_channels(64)_rows(2).ckpt ,然后导出对应的模型

from main import VPRModel
import torch
import onnxsim
import onnx
model = VPRModel(backbone_arch='resnet50',layers_to_crop=[4],agg_arch='MixVPR',agg_config={'in_channels': 1024,'in_h': 20,'in_w': 20,'out_channels': 64,'mix_depth': 4,'mlp_ratio': 1,'out_rows': 2},)state_dict = torch.load('resnet50_MixVPR_128_channels(64)_rows(2).ckpt')
model.load_state_dict(state_dict)
model.eval()
model.cpu()dummy_input = torch.randn(1, 3, 320, 320)
# 导出模型为ONNX格式
onnx_file_path = "vpr_model.onnx"
torch.onnx.export(model, dummy_input, # 模型输入的图片onnx_file_path, export_params=True, opset_version=12,  # 确保使用适合的ONNX opset版本do_constant_folding=True,  # 是否进行常量折叠优化input_names=['input'],  # 输入名称output_names=['output'],  # 输出名称dynamic_axes={'input': {0: 'batch_size'},  # 动态batch size'output': {0: 'batch_size'}})model_sim, flag = onnxsim.simplify(onnx_file_path)
if flag:onnx.save(model_sim, onnx_file_path)print("---------simplify onnx successfully---------")
else:print("---------simplify onnx failed-----------")

以上步骤中，我们完成了从 PyTorch 模型到 ONNX 格式的导出，并对 ONNX 模型进行了简化。简化后的 ONNX 模型不仅更加高效，还能够在部署时减少计算开销。

3. ONNX模块修改

这里值得一提的是，我们可以通过netron和onnx-graphsurgeon来完成模块的替换修改

在实际部署中，我们可能需要对模型的某些部分进行调整。例如，某些操作可能在某些硬件平台上表现不佳，或者我们需要替换特定的操作以实现更高效的推理。通过 Netron 可视化工具和 onnx-graphsurgeon 库，我们可以轻松地查看和修改 ONNX 模型中的节点。

以下是如何使用 onnx-graphsurgeon 替换一个 Sqrt 操作为 Log 操作的示例代码：

import onnx
import onnx_graphsurgeon as gs# 加载 ONNX 模型
onnx_model = onnx.load("vpr_model.onnx")# 将 ONNX 模型转换为 GraphSurgeon 的图
graph = gs.import_onnx(onnx_model)# 查找节点函数
def find_node_by_name(graph, node_name):for node in graph.nodes:if node.name == node_name:return nodereturn None# 查找 Sqrt 节点
sqrt_node = find_node_by_name(graph, "/aggregator/mix/mix.0/mix/mix.0/Sqrt")
sub_node = find_node_by_name(graph, "/aggregator/mix/mix.0/mix/mix.0/Sub")if sqrt_node:# 记录 `Sqrt` 节点的输入和输出sqrt_input = sqrt_node.inputs[0]  # 假设 Sqrt 节点有且仅有一个输入sqrt_output = sqrt_node.outputs[0]  # 假设 Sqrt 节点有且仅有一个输出print("Sqrt node found. Input:", sqrt_input, "Output:", sqrt_output)# 获取连接到 `Sqrt` 输出的节点（如 Log 节点）connected_nodes = [node for node in graph.nodes if sqrt_output in node.inputs]# 将 `Sqrt` 的输入直接连接到这些节点上for node in connected_nodes:node.inputs = [input if input != sqrt_output else sqrt_input for input in node.inputs]log_output = sqrt_output.copy()log_output.name = "new_Log_node"# 创建新的 Log 节点，连接 `Sqrt` 节点的输入和原来的输出new_log_node = gs.Node(op="Log", name="new_Log_node", inputs=[sqrt_input], outputs=[log_output])# sqrt_node.inputs[0] = "None"# 删除 `Sqrt` 节点graph.nodes.remove(sqrt_node)for node in graph.nodes:if node.name == "/aggregator/mix/mix.0/mix/mix.0/Div":node.inputs= [log_output,sub_node.outputs[0]]# 将新节点插入到图中breakgraph.nodes.append(new_log_node)# Cleanup 和重新排序graph.cleanup()graph.toposort()# # 打印修改后的图的节点# print("Graph after replacing the Sqrt node with Log node:")# for node in graph.nodes:#     print(node.name, node.op)# 将修改后的图导出为 ONNX 模型onnx_model_cut = gs.export_onnx(graph)onnx.save(onnx_model_cut, "vpr_model_cut.onnx")print("Model with Sqrt node replaced by Log node saved as vpr_model_cut.onnx")
else:print("Sqrt node could not be found, so the graph was not modified.")

通过这个例子，我们展示了如何在不影响模型整体结构的情况下，替换模型中的特定操作。这种灵活性使我们能够针对不同的硬件和应用需求，优化模型的性能。

4. TensorRT

为了加速模型推理，我们可以使用 NVIDIA 的 TensorRT。TensorRT 是一个用于深度学习推理优化的高性能库，它可以显著提高模型的推理速度。通过将 ONNX 模型转换为 TensorRT 引擎文件 (.engine)，我们可以在 NVIDIA GPU 上高效地运行推理任务。

首先，我们需要确保安装了合适版本的 TensorRT。这里我们可以根据我们之前从BEVDET来学习如何生成trt以及如何去写这些C++内容这一篇文章提到如何来下载安装TensorRT。在本次实验中，我们使用的是 TensorRT-8.6.1.6 版本。安装好 TensorRT 后，可以通过以下命令将 ONNX 模型转换为 TensorRT 引擎：

trtexec --onnx='/xxx/MixVPR/vpr_model.onnx' --fp16 --saveEngine=mix1.engine --warmUp=500 --duration=10 

该命令将 vpr_model.onnx 模型转换为 FP16 精度的 TensorRT 引擎文件 mix1.engine。其中，–warmUp=500 和 --duration=10 用于测试引擎性能和稳定性。

到这一步，我们就已经完成了整个模型的生成，并且获得了一个 mix1.engine 的文件。接下来，我们就可以编写 TensorRT 的 C++ 代码来进行推理了。

到这一步，我们就已经完成了整个模型的生成。并且获得了一个mix1.engine的文件。然后就可以写Tensorrt的代码了

4.1 main.cpp

main.cpp 是实现 MixVPR 模型推理过程的核心文件。MixVPR 类封装了使用 TensorRT 进行推理、读取图像数据、运行图像检索算法并将结果可视化的整个过程。主要步骤包括加载 TensorRT 引擎、执行推理、使用 FAISS 进行相似度搜索，以及展示匹配结果。

#include "opencv2/opencv.hpp"#include "cuda_runtime_api.h"
#include "opencv2/core.hpp"
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include "opencv2/highgui.hpp"#include "faiss/IndexFlat.h"
#include "faiss/IndexHNSW.h"#include "tensorrt_utils.h"
#include <iostream>
class frame {
public:
//    std::string frame_id;
//    cv::String filename;
//    cv::Mat raw_image;std::vector<float> img_global_des_vec;std::vector<float> local_descriptor;std::vector<float> kpoints;std::vector<float> similarity;std::vector<std::pair<int,double>> top_k_ind;std::vector<std::pair<float, float>> landmarks;
};using idx_t = faiss::idx_t;class MixVPR {
public:MixVPR(const std::string& engine_path);void inference(const cv::String& filename, std::vector<frame>& frame_set, std::vector<float>& des_db);void test_in_datasets(const std::string filepath, std::vector<cv::String>& namearray, std::vector<frame>& frame_set, std::vector<float>& des_db);void sort_vec_faiss(std::vector<frame>& frame_set, std::vector<float>& des_db);void run(std::string& datapath);private:std::shared_ptr<nvinfer1::ICudaEngine> engine;std::shared_ptr<nvinfer1::IExecutionContext> execution_context;TRTLogger Logger;
};MixVPR::MixVPR(const std::string& engine_path) {static auto engine_data = load_file(engine_path);static auto runtime = make_nvshared(nvinfer1::createInferRuntime(Logger));engine = make_nvshared(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));if (!engine) {printf("Deserialize cuda engine failed.\n");runtime->destroy();}execution_context = make_nvshared(engine->createExecutionContext());
}void MixVPR::inference(const cv::String& filename, std::vector<frame>& frame_set, std::vector<float>& des_db) {frame frame;cudaStream_t stream = nullptr;checkRuntime(cudaStreamCreate(&stream));int input_numel = 1 * 3 * 320 * 320;float* input_data_host = nullptr;float* input_data_device = nullptr;checkRuntime(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));cv::Mat image = cv::imread(filename);cv::resize(image, image, cv::Size(320, 320));unsigned char* pimage = image.data;for (int j = 0; j < image.cols * image.rows; ++j, pimage += 3) {input_data_host[j] = (pimage[2] / 255.0f - 0.485) / 0.229;input_data_host[j + image.cols * image.rows] = (pimage[1] / 255.0f - 0.456) / 0.224;input_data_host[j + 2 * image.cols * image.rows] = (pimage[0] / 255.0f - 0.406) / 0.225;}checkRuntime(cudaMemcpyAsync(input_data_device, input_data_host, input_numel * sizeof(float), cudaMemcpyHostToDevice, stream));float output_data_host[128];// 64*2float* output_data_device = nullptr;checkRuntime(cudaMalloc(&output_data_device, sizeof(output_data_host)));execution_context->setBindingDimensions(0, nvinfer1::Dims4{1, 3, 320, 320});float* bindings[] = {input_data_device, output_data_device};execution_context->enqueueV2((void**)bindings, stream, nullptr);checkRuntime(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream));checkRuntime(cudaStreamSynchronize(stream));frame.img_global_des_vec.insert(frame.img_global_des_vec.begin(), output_data_host, output_data_host + 128);des_db.insert(des_db.end(), output_data_host, output_data_host + 128);cudaStreamDestroy(stream);cudaFreeHost(input_data_host);cudaFree(input_data_device);cudaFree(output_data_device);frame_set.push_back(frame);
}void MixVPR::test_in_datasets(const std::string filepath, std::vector<cv::String>& namearray, std::vector<frame>& frame_set, std::vector<float>& des_db) {cv::glob(filepath, namearray);std::sort(namearray.begin(), namearray.end());for (size_t i = 0; i < namearray.size(); i++) {inference(namearray[i], frame_set, des_db);
//        frame_set[i].raw_image = cv::imread(namearray[i]);
//        frame_set[i].frame_id = std::to_string(i);
//        frame_set[i].filename = namearray[i];}
}void MixVPR::sort_vec_faiss(std::vector<frame>& frame_set, std::vector<float>& des_db) {faiss::IndexFlatIP index(128);index.add(frame_set.size(), des_db.data());//按照索引添加数据for (auto& img : frame_set) {idx_t I[3];float D[3];index.search(1, img.img_global_des_vec.data(), 3, D, I);img.top_k_ind.push_back(std::make_pair(I[1], D[1]));}
}void MixVPR::run(std::string& datapath) {std::vector<cv::String> namearray;std::vector<frame> frame_set;std::vector<float> des_db;test_in_datasets(datapath, namearray, frame_set, des_db);sort_vec_faiss(frame_set, des_db);for (const auto& img : frame_set) {auto img0 = cv::imread(namearray[img.top_k_ind[0].first]);//输出结果auto img1 = cv::imread(namearray[&img - &frame_set[0]]);//对应图片cv::Mat img2;if (img0.empty() || img1.empty()) {std::cout << "Read image failed." << std::endl;continue;}cv::resize(img0, img0, cv::Size(720, 720));cv::resize(img1, img1, cv::Size(720, 720));if (img.top_k_ind[0].second > 0.75) {cv::hconcat(img1, img0, img2);cv::putText(img2, "Matched", cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 2);} else {cv::Mat empty_img = cv::Mat::zeros(img0.size(), img0.type());cv::hconcat(img1, empty_img, img2);cv::putText(img2, "Unmatched", cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 2);}cv::imshow("Result", img2);cv::waitKey(0);}
}int main(int argc, char** argv) {std::string datapath = "/xxx/MixVPR/Tensor/pic";MixVPR mixvpr("/xxxx/MixVPR/mix1.engine");mixvpr.run(datapath);return 0;
}

4.2 tensorrt_utils.cpp

这个文件实现了一些帮助函数和日志记录器，用于支持 TensorRT 推理流程。其中包括检查 CUDA 运行时错误的函数、日志记录器的实现（用于记录 TensorRT 的输出信息），以及从文件加载数据的函数。这些工具函数简化了在主程序中与 TensorRT 和 CUDA 进行交互的代码。

#include "tensorrt_utils.h"using namespace std;bool __check_cuda_runtime(cudaError_t code, const char* op, const char* file, int line){if(code != cudaSuccess){const char* err_name = cudaGetErrorName(code);const char* err_message = cudaGetErrorString(code);printf("runtime error %s:%d %s failed. \n code = %s, message =%s\n", file, line, op, err_name, err_message);return false;}return true;}inline const char* severity_string(nvinfer1::ILogger::Severity t){switch(t){case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";case nvinfer1::ILogger::Severity::kERROR: return "error";case nvinfer1::ILogger::Severity::kWARNING: return "warning";case nvinfer1::ILogger::Severity::kINFO: return "info";case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";default: return "unknow";}
}void TRTLogger::log(nvinfer1::ILogger::Severity severity, const nvinfer1::AsciiChar *msg) noexcept {if(severity <= Severity::kINFO){if(severity == Severity::kWARNING){printf("\033[33m%s: %s\033[0m\n", severity_string(severity),msg);}else if(severity <= Severity::kERROR){printf("\033[31m%s: %s\033[0m\n", severity_string(severity),msg);}else{printf("%s: %s\n", severity_string(severity), msg);}}
}vector<unsigned char> load_file(const string& file){ifstream in(file, ios::in | ios::binary);if (!in.is_open())return {};in.seekg(0, ios::end);size_t length = in.tellg();std::vector<uint8_t> data;if (length > 0){in.seekg(0, ios::beg);data.resize(length);in.read((char*)&data[0], length);}in.close();return data;}

4.3 tensorrt_utils.h

这是 tensorrt_utils.cpp 的头文件，声明了该文件中的函数和类。它定义了 TensorRT 日志记录器类 TRTLogger 以及一些帮助函数和宏（例如 checkRuntime 用于 CUDA 错误检查）。这些声明为主程序提供了所需的接口和工具函数支持

#ifndef VINS_FUSION_TENSORRT_UTILS_H
#define VINS_FUSION_TENSORRT_UTILS_H#pragma once
//tensorrt
#include "NvInfer.h"
#include "NvInferRuntime.h"//cuda
#include "cuda.h"
#include <cuda_runtime.h>//sys
#include "vector"
#include "fstream"
#include <memory>//tensorrt_tools
// #include "tensorrt_tools/trt_infer.hpp"
// #include "tensorrt_tools/trt_tensor.hpp"
// #include "tensorrt_tools/cuda_tools.hpp"
// #include "tensorrt_tools/ilogger.hpp"#define checkRuntime(op) __check_cuda_runtime((op), #op, __FILE__ ,__LINE__)using namespace std;bool __check_cuda_runtime(cudaError_t code, const char* op, const char* file, int line);inline const char* severity_string(nvinfer1::ILogger::Severity t);class TRTLogger : public nvinfer1::ILogger{
public:virtual void log(Severity severity, nvinfer1::AsciiChar const* msg)noexcept override;
};template<typename _T>
static shared_ptr<_T> make_nvshared(_T* ptr){//通过智能指针管理 nv 返回的指针参数，内存自动释放，避免泄漏return shared_ptr<_T>(ptr, [](_T* p){p->destroy();});
}vector<unsigned char> load_file(const string& file);
#endif //VINS_FUSION_TENSORRT_UTILS_H

4.4 CMakeLists.txt

cmake_minimum_required(VERSION 3.10)
project(MixVPR)set(CMAKE_CXX_STANDARD 14)# 找到 OpenCV
find_package(OpenCV REQUIRED)# 找到 CUDA
find_package(CUDA REQUIRED)
find_package(faiss REQUIRED)# 设置 TensorRT 的路径
set(TENSORRT_ROOT /xxxx/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8/TensorRT-8.6.1.6)
find_path(TENSORRT_INCLUDE_DIR NvInfer.hHINTS ${TENSORRT_ROOT}/include)
find_library(TENSORRT_LIBRARY nvinferHINTS ${TENSORRT_ROOT}/lib)
find_library(TENSORRT_RUNTIME_LIBRARY nvparsersHINTS ${TENSORRT_ROOT}/lib)# 包含 TensorRT 的头文件路径
include_directories(${TENSORRT_INCLUDE_DIR})# 包含 OpenCV 和 CUDA 的头文件路径
include_directories(${OpenCV_INCLUDE_DIRS})
include_directories(${CUDA_INCLUDE_DIRS})# 编译项目
add_executable(MixVPR main.cpp tensorrt_utils.cpp)# 链接 TensorRT, OpenCV 和 CUDA 的库
target_link_libraries(MixVPR ${TENSORRT_LIBRARY} ${TENSORRT_RUNTIME_LIBRARY} ${OpenCV_LIBS} ${CUDA_LIBRARIES} faiss)