YOLOv8解码 tensorrt

### 使用TensorRT解码YOLOv8模型 #### 准备工作 为了使用TensorRT对YOLOv8模型进行解码，首先需要准备环境并加载必要的库文件。确保已安装CUDA、cuDNN和TensorRT，并配置好相应的环境变量。 ```bash export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/TensorRT/lib ``` #### 加载ONNX模型转换为TensorRT引擎 通常情况下，YOLOv8的训练是在PyTorch框架内完成的，之后可以导出为ONNX格式以便于后续优化处理。通过`trtexec`工具或者编写自定义程序来读取该ONNX文件并将其转化为高效的TensorRT推理引擎。 ```cpp #include "NvInfer.h" // ...其他头文件... nvinfer1::ICudaEngine* create_engine(int batch_size, const std::string& onnx_file_path){ // 创建Builder对象用于构建网络结构 nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(gLogger); // 获取显卡信息以确定最大批处理量和支持特性 int nbGpuDevices = getNbAvailableGpus(); builder->setMaxBatchSize(batch_size); // 设置精度模式 FP16 或者 INT8 如果硬件支持的话 builder->setMinOptimizationLevel(3); // 解析ONNX模型创建INetworkDefinition实例 ONNXParser::IParser* parser = nvonnxparser::createParser(*network, gLogger); if (!parser->parseFromFile(onnx_file_path.c_str(), static_cast<int>(ILogger::Severity::kINFO))) { throw std::runtime_error("Failed to parse the ONNX file"); } // 构建序列化后的engine保存到磁盘上供以后快速加载 nvinfer1::IHostMemory* modelStream{nullptr}; auto engine = builder->buildCudaEngine(*network); } ``` #### 执行推理过程中的解码操作 对于YOLOv8而言，在得到预测结果后还需要进一步解析这些数据才能获取最终的目标检测框位置与类别概率分布。这部分逻辑涉及到非极大抑制(NMS)，边界框坐标变换等内容。 ```cpp void decode_outputs(const float* output_data, size_t num_boxes, size_t num_classes, std::vector<Detection>& detections) { constexpr float CONFIDENCE_THRESHOLD = 0.5f; constexpr float NMS_IOU_THRESHOLD = 0.4f; for (size_t i = 0; i < num_boxes; ++i) { const float objectness_score = sigmoid(output_data[i * LEN_ALL_RESULT + NUM_CLASSES]); // 计算置信度得分[^3] if (objectness_score >= CONFIDENCE_THRESHOLD) { // 过滤掉低分候选框 BoundingBox bbox{}; // 提取出分类分数最高的那一类作为当前box所属标签 float max_class_confidence = -FLT_MAX; int best_class_idx = -1; for (int c = 0; c < NUM_CLASSES; ++c) { float class_confidence = sigmoid(output_data[i * LEN_ALL_RESULT + c]); if (class_confidence > max_class_confidence) { max_class_confidence = class_confidence; best_class_idx = c; } } // 将相对偏移量映射回原始图像空间内的绝对坐标系下表示形式 bbox.x_center = (output_data[i * LEN_ALL_RESULT + NUM_CLASSES + 0] / grid_width_) * input_image_width_; bbox.y_center = (output_data[i * LEN_ALL_RESULT + NUM_CLASSES + 1] / grid_height_) * input_image_height_; bbox.width = exp(output_data[i * LEN_ALL_RESULT + NUM_CLASSES + 2]) * anchor_w_[best_anchor_index_]; bbox.height = exp(output_data[i * LEN_ALL_RESULT + NUM_CLASSES + 3]) * anchor_h_[best_anchor_index_]; Detection det{ .bbox = bbox, .confidence = objectness_score * max_class_confidence, .label = labels_.at(best_class_idx), }; detections.push_back(det); } } // 应用NMS去除冗余重叠较高的多个相同物体提议 apply_nms(detections, NMS_IOU_THRESHOLD); } std::vector<std::string> labels_{"person", /*...*/}; // 定义所有可能存在的目标名称列表 float anchor_w_[9]{/*预设先验框宽度数组*/}, anchor_h_[9]; // YOLO使用的锚点尺寸参数表 unsigned int grid_width_, grid_height_; // 特征图网格大小 unsigned int input_image_width_, input_image_height_; // 输入图片分辨率规格 ``` 上述代码片段展示了如何基于TensorRT运行时API实现YOLOv8模型的具体解码流程。需要注意的是实际应用中还需考虑更多细节如批量处理能力调整、多尺度测试策略等。