tensorrtx yolov8

### 使用 TensorRTX 构建 YOLOv8 的 DLL 文件 为了使用 TensorRTX 构建 YOLOv8 的动态链接库 (DLL)，需要遵循一系列特定的操作流程。以下是详细的说明： #### 准备环境 确保安装并配置好 NVIDIA CUDA 和 cuDNN，以及最新版本的 TensorRT SDK[^2]。 #### 下载和准备源码 获取最新的 `tensorrtx` 仓库，并下载预训练好的 YOLOv8 权重文件 `.pt` 或者转换成适合 TensorRT 输入格式如 ONNX 格式的模型文件[^1]。 #### 转换模型至ONNX格式 如果原始权重不是ONNX格式，则需先将其转为该格式以便后续处理： ```python import torch from yolov5.models.experimental import attempt_load device = 'cuda' if torch.cuda.is_available() else 'cpu' model = attempt_load('path_to_yolov8.pt', map_location=device) # 加载PyTorch模型 dummy_input = torch.randn(1, 3, 640, 640).to(device) torch.onnx.export(model, dummy_input, "yolov8.onnx", opset_version=11) ``` #### 修改 CMakeLists.txt 配置编译选项 编辑 `tensorrtx/yolov5/CMakeLists.txt` 文件以适应新的架构需求，特别是针对 YOLOv8 特定层的支持。 #### 编写导出函数定义头文件 (.h) 创建一个新的头文件用于声明外部可见的功能接口，例如 `export_yolov8.h`: ```cpp #ifndef EXPORT_YOLOV8_H_ #define EXPORT_YOLOV8_H_ #ifdef __cplusplus extern "C" { #endif // 导出检测功能 __declspec(dllexport) void* createYoloDetector(const char*); __declspec(dllexport) int detectObjects(void*, float***, int**, float); __declspec(dllexport) void destroyYoloDetector(void*); #ifdef __cplusplus } #endif #endif //EXPORT_YOLOV8_H_ ``` #### 实现对应的 CPP 源代码实现文件(.cpp) 编写实际逻辑来加载 ONNX 模型并通过 TensorRT 进行推理，保存为 `export_yolov8.cpp`: ```cpp #include <NvInfer.h> #include "export_yolov8.h" static ILogger gLogger; void* detector; // 创建探测器实例... void* createYoloDetector(const char *onnxFile){ // ...省略具体细节... } int detectObjects(/* 参数列表 */){ // 执行推断操作... } void destroyYoloDetector(){ // 清理资源... } ``` #### 编译生成 DLL 文件 利用上述修改后的项目结构，在命令提示符窗口中执行如下指令完成最终目标——生成可供调用的共享库(`.dll`)文件: ```bash mkdir build && cd build cmake .. cmake --build . --config Release ``` 成功之后会在 `build/Release` 目录找到所需的 DLL 输出产物。

### 使用 TensorRTX 优化和部署 YOLOv5 模型的流程 在使用 TensorRTX 优化和部署 YOLOv5 模型时，需要完成多个关键步骤，包括环境准备、模型转换以及最终的推理部署。以下是详细说明： #### 1. 环境准备 确保安装了必要的依赖库和工具链。例如，TensorRTX 支持多种操作系统（如 Linux 和 Windows），但推荐在 Ubuntu 系统上进行开发和部署[^2]。此外，下载与 YOLOv5 版本匹配的 TensorRTX 代码库。例如，如果使用的是 YOLOv5 v7.0，则应下载对应的 TensorRTX 分支[^3]。 ```bash git clone -b v7.0 https://github.com/ultralytics/yolov5.git git clone -b yolov5-v7.0 https://github.com/wang-xinyu/tensorrtx.git ``` #### 2. 模型转换 将训练好的 PyTorch (.pt) 模型转换为 TensorRT 支持的格式（如 .wts 或 .onnx）。首先，利用 `gen_wts.py` 脚本将 `.pt` 文件转换为 `.wts` 格式。此脚本通常位于 TensorRTX 的 YOLOv5 工程目录下[^1]。 ```python # 将 gen_wts.py 拷贝到 YOLOv5 工程目录并运行 cp tensorrtx/yolov5/gen_wts.py ultralytics/yolov5/ cd ultralytics/yolov5 python gen_wts.py -w yolov5s.pt -o yolov5s.wts ``` 接下来，可以将 `.wts` 文件进一步转换为 ONNX 格式[^4]，这是 TensorRT 推理引擎支持的标准中间表示。 ```python import torch from models.experimental import attempt_load # 加载 .pt 模型 model = attempt_load('yolov5s.pt', map_location=torch.device('cpu')) # 导出为 ONNX 格式 torch.onnx.export( model, torch.zeros(1, 3, 640, 640), 'yolov5s.onnx', input_names=['images'], output_names=['output'] ) ``` #### 3. 构建 TensorRT 引擎 使用 TensorRTX 提供的 C++ 或 Python 工具生成 TensorRT 引擎文件（`.engine`）。以下是一个简单的 C++ 示例，展示如何加载 ONNX 模型并生成 TensorRT 引擎[^5]。 ```cpp #include "NvInfer.h" #include <iostream> #include <fstream> using namespace nvinfer1; IBuilder* builder = createInferBuilder(gLogger); INetworkDefinition* network = builder->createNetworkV2(0U); // 解析 ONNX 模型 onnx::Parser parser(network); if (!parser.parseFromFile("yolov5s.onnx", static_cast<int>(gLogger.getSeverity()))) { std::cerr << "Failed to parse ONNX file" << std::endl; return -1; } // 配置构建器 IBuilderConfig* config = builder->createBuilderConfig(); config->setMaxWorkspaceSize(1 << 30); // 设置最大工作区大小 // 构建引擎 ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); if (!engine) { std::cerr << "Failed to build engine" << std::endl; return -1; } // 保存引擎 std::ofstream engineFile("yolov5s.engine", std::ios::binary); if (!engineFile) { std::cerr << "Failed to open engine file for writing" << std::endl; return -1; } engineFile.write(reinterpret_cast<const char*>(engine->serialize()), engine->getSerializedSize()); ``` #### 4. 推理部署 最后，在目标设备上加载生成的 `.engine` 文件，并执行推理任务。可以通过 TensorRT 提供的 API 实现高效的推理操作[^2]。 ```python import tensorrt as trt # 初始化 TensorRT 运行时 runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING)) # 加载引擎文件 with open("yolov5s.engine", "rb") as f: engine = runtime.deserialize_cuda_engine(f.read()) # 创建上下文 context = engine.create_execution_context() # 准备输入输出缓冲区 input_buffer = ... # 输入数据 output_buffer = ... # 输出缓冲区 # 执行推理 context.execute_v2([input_buffer, output_buffer]) ``` --- ###