### 加载 YOLOv 模型并获取输出张量维度的实现 在 C++ 中使用 TensorRT 加载 YOLOv 模型（如 YOLOv3 或 YOLOv8），需要通过解析 ONNX 文件生成 TensorRT 引擎，并获取输出张量的维度。以下是具体的实现方法： #### 1. 初始化 TensorRT 环境 首先，需要创建一个 `nvinfer1::IBuilder` 对象来构建 TensorRT 引擎。同时，创建一个 `nvinfer1::IRuntime` 对象用于加载序列化的引擎文件。 ```cpp #include <NvInfer.h> #include <fstream> #include <iostream> // 创建 TensorRT 运行时 nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger); ``` #### 2. 从 ONNX 文件构建 TensorRT 引擎 使用 `onnxParser` 解析 ONNX 文件并生成 TensorRT 引擎。 ```cpp // 读取 ONNX 文件到内存 std::ifstream file("yolov8n-obb.onnx", std::ios::binary); if (!file) { std::cerr << "Error: could not open file." << std::endl; return -1; } 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); // 创建解析器 nvinfer1::IHostMemory* modelStream{nullptr}; nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(buffer.data(), size, nullptr); if (!engine) { std::cerr << "Failed to create engine" << std::endl; return -1; } ``` #### 3. 获取输出张量的维度 通过 `nvinfer1::ICudaEngine` 的接口可以获取所有输出张量的名称和维度信息。 ```cpp // 获取输出张量的数量 int nbOutputs = engine->getNbBindings(); for (int i = 0; i < nbOutputs; ++i) { // 判断是否为输出张量 if (engine->bindingIsInput(i)) continue; // 获取张量名称 const char* tensorName = engine->getBindingName(i); // 获取张量维度 nvinfer1::Dims dims = engine->getBindingDimensions(i); // 打印维度信息 std::cout << "Output Tensor Name: " << tensorName << std::endl; std::cout << "Dimensions: "; for (int j = 0; j < dims.nbDims; ++j) { std::cout << dims.d[j] << " "; } std::cout << std::endl; } ``` 以上代码展示了如何从 ONNX 文件加载模型、构建 TensorRT 引擎以及获取输出张量的维度[^1]。 #### 4. 配置推理参数 根据提供的 `config.txt` 文件内容，可以通过读取配置文件设置推理所需的参数，例如置信度阈值、IoU 阈值、输入图像大小等[^2]。 ```cpp #include <unordered_map> #include <string> #include <sstream> std::unordered_map<std::string, std::string> config; void parseConfig(const std::string& filename) { std::ifstream file(filename); if (!file.is_open()) { std::cerr << "Failed to open config file!" << std::endl; return; } std::string line; while (std::getline(file, line)) { if (line.empty() || line[0] == '#') continue; std::istringstream ss(line); std::string key, value; if (std::getline(ss, key, '=') && std::getline(ss, value)) { config[key] = value; } } } // 示例：读取配置文件 parseConfig("config.txt"); float rectConfidenceThreshold = std::stof(config["rectConfidenceThreshold"]); float iouThreshold = std::stof(config["iouThreshold"]); int imageSize = std::stoi(config["imageSize"]); ``` #### 注意事项 - 在实际应用中，确保 ONNX 模型与 TensorRT 兼容。 - 如果模型包含动态形状支持，则需要额外处理动态维度的情况。 - 输出张量的维度可能因模型架构不同而有所差异，需根据具体模型调整解析逻辑。 ```cpp // 动态形状支持示例 if (engine->hasImplicitBatchDimension()) { std::cout << "Model uses implicit batch dimension." << std::endl; } else { std::cout << "Model uses explicit batch dimension." << std::endl; } ```