sladro
e13cb3659c
- 创建基本项目结构和目录 - 添加CMake构建系统 - 实现基础的配置解析功能 - 添加YOLO推理框架支持 - 集成RTSP和视频流处理功能 - 添加性能监控和日志系统
577 lines
19 KiB
C++
577 lines
19 KiB
C++
#include "trt_inference.hpp"
|
||
#include "cuda_helper.hpp"
|
||
#include <fstream>
|
||
#include <NvOnnxParser.h>
|
||
#include <iostream>
|
||
#include <filesystem>
|
||
#include "../common/logger.hpp"
|
||
#include <opencv2/imgproc.hpp>
|
||
|
||
namespace pipeline {
|
||
|
||
TrtInference::TrtInference(const InferenceConfig& config)
|
||
: config_(config),
|
||
logger_(std::make_unique<detail::Logger>()),
|
||
stream_(std::make_unique<detail::CudaStream>()) {
|
||
|
||
// 创建预处理器配置
|
||
Preprocessor::Config preprocess_config;
|
||
preprocess_config.input_shape = config.model.input_shape;
|
||
preprocess_config.use_cuda = true; // 默认使用CUDA加速
|
||
preprocess_config.max_batch_size = config.max_batch_size;
|
||
preprocess_config.scale = 1.0f/255.0f; // 归一化系数
|
||
|
||
// 创建预处理器实例
|
||
preprocessor_ = std::make_unique<Preprocessor>(preprocess_config);
|
||
}
|
||
|
||
TrtInference::~TrtInference() {
|
||
destroy();
|
||
}
|
||
|
||
bool TrtInference::convertOnnxToEngine(const std::string& onnx_path, const std::string& engine_path) {
|
||
try {
|
||
Logger::info("=== Starting ONNX to TensorRT Conversion ===");
|
||
Logger::info("ONNX file: " + onnx_path);
|
||
|
||
// 检查文件是否存在
|
||
if (!std::filesystem::exists(onnx_path)) {
|
||
Logger::error("ONNX file does not exist: " + onnx_path);
|
||
return false;
|
||
}
|
||
|
||
// 创建 builder
|
||
std::unique_ptr<nvinfer1::IBuilder, void(*)(nvinfer1::IBuilder*)> builder(
|
||
nvinfer1::createInferBuilder(*logger_),
|
||
[](nvinfer1::IBuilder* b) { delete b; });
|
||
if (!builder) {
|
||
Logger::error("Failed to create TensorRT Builder");
|
||
return false;
|
||
}
|
||
|
||
// 创建网络定义
|
||
#pragma GCC diagnostic push
|
||
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
|
||
const auto explicit_batch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
|
||
#pragma GCC diagnostic pop
|
||
std::unique_ptr<nvinfer1::INetworkDefinition, void(*)(nvinfer1::INetworkDefinition*)> network(
|
||
builder->createNetworkV2(explicit_batch),
|
||
[](nvinfer1::INetworkDefinition* n) { delete n; });
|
||
if (!network) {
|
||
Logger::error("Failed to create network definition");
|
||
return false;
|
||
}
|
||
|
||
// 创建 ONNX 解析器
|
||
std::unique_ptr<nvonnxparser::IParser, void(*)(nvonnxparser::IParser*)> parser(
|
||
nvonnxparser::createParser(*network, *logger_),
|
||
[](nvonnxparser::IParser* p) { delete p; });
|
||
if (!parser) {
|
||
Logger::error("Failed to create ONNX parser");
|
||
return false;
|
||
}
|
||
|
||
// 解析 ONNX 文件
|
||
if (!parser->parseFromFile(onnx_path.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kINFO))) {
|
||
Logger::error("Failed to parse ONNX file");
|
||
return false;
|
||
}
|
||
|
||
// 创建构建配置
|
||
std::unique_ptr<nvinfer1::IBuilderConfig, void(*)(nvinfer1::IBuilderConfig*)> config(
|
||
builder->createBuilderConfig(),
|
||
[](nvinfer1::IBuilderConfig* c) { delete c; });
|
||
if (!config) {
|
||
Logger::error("Failed to create builder config");
|
||
return false;
|
||
}
|
||
|
||
// 修改优化配置部分
|
||
auto profile = builder->createOptimizationProfile();
|
||
auto input = network->getInput(0);
|
||
if (!input) {
|
||
Logger::error("Failed to get network input");
|
||
return false;
|
||
}
|
||
|
||
// 获取输入名称和维度
|
||
const char* input_name = input->getName();
|
||
nvinfer1::Dims dims = input->getDimensions();
|
||
|
||
// 确保所有维度都是正数
|
||
for (int i = 0; i < dims.nbDims; i++) {
|
||
if (dims.d[i] < 0) {
|
||
if (i == 0) {
|
||
// batch维度
|
||
dims.d[i] = 1;
|
||
} else {
|
||
// 其他维度使用配置中的值
|
||
dims.d[i] = config_.model.input_shape[i-1];
|
||
}
|
||
}
|
||
}
|
||
|
||
// 设置最小、最优和最大维度
|
||
nvinfer1::Dims min_dims = dims;
|
||
nvinfer1::Dims opt_dims = dims;
|
||
nvinfer1::Dims max_dims = dims;
|
||
|
||
// 只修改batch维度
|
||
min_dims.d[0] = 1; // 最小batch size
|
||
opt_dims.d[0] = 1; // 最优batch size
|
||
max_dims.d[0] = config_.max_batch_size; // 最大batch size
|
||
|
||
// 设置优化配置
|
||
profile->setDimensions(input_name, nvinfer1::OptProfileSelector::kMIN, min_dims);
|
||
profile->setDimensions(input_name, nvinfer1::OptProfileSelector::kOPT, opt_dims);
|
||
profile->setDimensions(input_name, nvinfer1::OptProfileSelector::kMAX, max_dims);
|
||
|
||
if (!profile->isValid()) {
|
||
Logger::error("Invalid optimization profile");
|
||
return false;
|
||
}
|
||
|
||
config->addOptimizationProfile(profile);
|
||
|
||
// 设置工作空间大小
|
||
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, config_.workspace_size);
|
||
|
||
// 设置精度
|
||
#pragma GCC diagnostic push
|
||
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
|
||
if (builder->platformHasFastFp16()) {
|
||
Logger::info("Enabling FP16 precision");
|
||
config->setFlag(nvinfer1::BuilderFlag::kFP16);
|
||
}
|
||
#pragma GCC diagnostic pop
|
||
|
||
// 构建引擎
|
||
std::unique_ptr<nvinfer1::IHostMemory, void(*)(nvinfer1::IHostMemory*)> serializedEngine(
|
||
builder->buildSerializedNetwork(*network, *config),
|
||
[](nvinfer1::IHostMemory* m) { delete m; });
|
||
if (!serializedEngine) {
|
||
Logger::error("Failed to build TensorRT engine");
|
||
return false;
|
||
}
|
||
|
||
// 保存引擎文件
|
||
std::ofstream engineFile(engine_path, std::ios::binary);
|
||
if (!engineFile) {
|
||
Logger::error("Failed to open engine file for writing");
|
||
return false;
|
||
}
|
||
|
||
engineFile.write(static_cast<const char*>(serializedEngine->data()), serializedEngine->size());
|
||
engineFile.close();
|
||
|
||
Logger::info("Successfully saved TensorRT engine to: " + engine_path);
|
||
return true;
|
||
}
|
||
catch (const std::exception& e) {
|
||
Logger::error("Error in convertOnnxToEngine: " + std::string(e.what()));
|
||
return false;
|
||
}
|
||
}
|
||
|
||
bool TrtInference::loadEngine() {
|
||
try {
|
||
// 确保CUDA上下文初始化
|
||
cudaFree(0); // 强制初始化CUDA上下文
|
||
|
||
// 检查文件是否存在
|
||
if (!std::filesystem::exists(config_.model.engine_path)) {
|
||
Logger::error("Engine file not found: " + config_.model.engine_path);
|
||
return false;
|
||
}
|
||
|
||
// 读取引擎文件
|
||
std::ifstream file(config_.model.engine_path, std::ios::binary);
|
||
if (!file) {
|
||
Logger::error("Failed to open engine file");
|
||
return false;
|
||
}
|
||
|
||
file.seekg(0, std::ios::end);
|
||
size_t size = file.tellg();
|
||
file.seekg(0, std::ios::beg);
|
||
|
||
std::vector<char> engineData(size);
|
||
file.read(engineData.data(), size);
|
||
file.close();
|
||
|
||
// 创建 runtime
|
||
runtime_ = nvinfer1::createInferRuntime(*logger_);
|
||
if (!runtime_) {
|
||
Logger::error("Failed to create TensorRT Runtime");
|
||
return false;
|
||
}
|
||
|
||
// 反序列化引擎
|
||
engine_ = runtime_->deserializeCudaEngine(engineData.data(), size);
|
||
if (!engine_) {
|
||
Logger::error("Failed to deserialize TensorRT engine");
|
||
return false;
|
||
}
|
||
|
||
// 创建执行上下文
|
||
if (!createContext()) {
|
||
Logger::error("Failed to create execution context");
|
||
return false;
|
||
}
|
||
|
||
// 分配缓冲区
|
||
if (!allocateBuffers()) {
|
||
Logger::error("Failed to allocate CUDA buffers");
|
||
return false;
|
||
}
|
||
|
||
// 同步设备
|
||
cudaDeviceSynchronize();
|
||
|
||
Logger::info("Successfully loaded TensorRT engine");
|
||
return true;
|
||
}
|
||
catch (const std::exception& e) {
|
||
Logger::error("Error in loadEngine: " + std::string(e.what()));
|
||
return false;
|
||
}
|
||
}
|
||
|
||
bool TrtInference::infer(const std::vector<cv::Mat>& images, std::vector<DetectionResult>& results) {
|
||
if (!isLoaded()) {
|
||
Logger::error("Engine not loaded");
|
||
return false;
|
||
}
|
||
|
||
try {
|
||
// 确保CUDA上下文有效
|
||
cudaFree(0);
|
||
|
||
// 1. 预处理
|
||
if (!preprocess(images)) {
|
||
return false;
|
||
}
|
||
|
||
// 2. 设置输入形状
|
||
if (!context_->setInputShape("images",
|
||
nvinfer1::Dims4(images.size(), config_.model.input_shape[0],
|
||
config_.model.input_shape[1], config_.model.input_shape[2]))) {
|
||
Logger::error("Failed to set input shape");
|
||
return false;
|
||
}
|
||
|
||
// 3. 设置输入和输出张量地址
|
||
if (!context_->setTensorAddress("images", input_buffers_[0]->devicePtr())) {
|
||
Logger::error("Failed to set input tensor address");
|
||
return false;
|
||
}
|
||
|
||
const char* output_name = engine_->getIOTensorName(engine_->getNbIOTensors() - 1);
|
||
if (!context_->setTensorAddress(output_name, output_buffers_[0]->devicePtr())) {
|
||
Logger::error("Failed to set output tensor address");
|
||
return false;
|
||
}
|
||
|
||
// 4. 执行推理
|
||
if (!context_->enqueueV3(stream_->get())) {
|
||
Logger::error("Inference failed");
|
||
return false;
|
||
}
|
||
|
||
// 5. 后处理
|
||
if (!postprocess(results)) {
|
||
return false;
|
||
}
|
||
|
||
// 6. 同步流和设备
|
||
if (!stream_->sync()) {
|
||
Logger::error("Failed to synchronize CUDA stream");
|
||
return false;
|
||
}
|
||
cudaDeviceSynchronize();
|
||
|
||
return true;
|
||
}
|
||
catch (const std::exception& e) {
|
||
Logger::error("Error in inference: " + std::string(e.what()));
|
||
return false;
|
||
}
|
||
}
|
||
|
||
bool TrtInference::preprocess(const std::vector<cv::Mat>& images) {
|
||
try {
|
||
// 检查输入
|
||
if (images.empty() || images.size() > static_cast<size_t>(config_.max_batch_size)) {
|
||
Logger::error("Invalid batch size: " + std::to_string(images.size()));
|
||
return false;
|
||
}
|
||
|
||
// 使用预处理器处理图像
|
||
if (!preprocessor_->process(images, *input_buffers_[0])) {
|
||
Logger::error("Preprocessing failed");
|
||
return false;
|
||
}
|
||
|
||
return true;
|
||
}
|
||
catch (const std::exception& e) {
|
||
Logger::error("Error in preprocessing: " + std::string(e.what()));
|
||
return false;
|
||
}
|
||
}
|
||
|
||
bool TrtInference::postprocess(std::vector<DetectionResult>& results) {
|
||
try {
|
||
// 获取输出尺寸
|
||
const auto output_dims = engine_->getTensorShape(output_names_[0].c_str());
|
||
const int num_outputs = output_dims.d[1]; // 每个检测的输出维度
|
||
const int max_detections = output_dims.d[2]; // 最大检测数量
|
||
|
||
// 分配主机内存接收结果
|
||
const size_t output_size = output_buffers_[0]->size();
|
||
std::vector<float> output_data(output_size / sizeof(float));
|
||
|
||
// 从GPU拷贝结果
|
||
if (cudaMemcpyAsync(output_data.data(), output_buffers_[0]->devicePtr(),
|
||
output_size, cudaMemcpyDeviceToHost, stream_->get()) != cudaSuccess) {
|
||
Logger::error("Failed to copy output data from GPU");
|
||
return false;
|
||
}
|
||
|
||
// 等待拷贝完成
|
||
if (cudaStreamSynchronize(stream_->get()) != cudaSuccess) {
|
||
Logger::error("Failed to synchronize CUDA stream");
|
||
return false;
|
||
}
|
||
|
||
// 解析结果
|
||
results.clear();
|
||
results.resize(1); // 每批次的结果
|
||
auto& result = results[0];
|
||
|
||
// YOLOv8输出格式：[batch_id, x, y, w, h, conf, class_conf, class_id]
|
||
for (int i = 0; i < max_detections; ++i) {
|
||
const float* det = output_data.data() + i * num_outputs;
|
||
const float conf = det[5];
|
||
|
||
if (conf < config_.threshold.conf) continue;
|
||
|
||
BBox box;
|
||
box.x1 = det[1] - det[3] / 2.0f; // center_x - width/2
|
||
box.y1 = det[2] - det[4] / 2.0f; // center_y - height/2
|
||
box.x2 = det[1] + det[3] / 2.0f; // center_x + width/2
|
||
box.y2 = det[2] + det[4] / 2.0f; // center_y + height/2
|
||
box.score = conf; // 使用 score 而不是 conf
|
||
box.class_id = static_cast<int>(det[7]);
|
||
|
||
result.boxes.push_back(box);
|
||
}
|
||
|
||
return true;
|
||
}
|
||
catch (const std::exception& e) {
|
||
Logger::error("Error in postprocessing: " + std::string(e.what()));
|
||
return false;
|
||
}
|
||
}
|
||
|
||
void TrtInference::doNMS(std::vector<BBox>& boxes, float nms_thresh) {
|
||
// 按置信度排序，修改 lambda 表达式确保返回 bool 值
|
||
std::sort(boxes.begin(), boxes.end(),
|
||
[](const BBox& a, const BBox& b) -> bool {
|
||
return a.score > b.score; // 使用 score 而不是 conf
|
||
});
|
||
|
||
std::vector<bool> keep(boxes.size(), true);
|
||
|
||
// NMS
|
||
for (size_t i = 0; i < boxes.size(); ++i) {
|
||
if (!keep[i]) continue;
|
||
|
||
for (size_t j = i + 1; j < boxes.size(); ++j) {
|
||
if (!keep[j]) continue;
|
||
|
||
// 如果是同一类别且IOU大于阈值，则抑制
|
||
if (boxes[i].class_id == boxes[j].class_id &&
|
||
calculateIOU(boxes[i], boxes[j]) > nms_thresh) {
|
||
keep[j] = false;
|
||
}
|
||
}
|
||
}
|
||
|
||
// 移除被抑制的框
|
||
std::vector<BBox> nms_boxes;
|
||
for (size_t i = 0; i < boxes.size(); ++i) {
|
||
if (keep[i]) {
|
||
nms_boxes.push_back(boxes[i]);
|
||
}
|
||
}
|
||
|
||
boxes.swap(nms_boxes);
|
||
}
|
||
|
||
float TrtInference::calculateIOU(const BBox& box1, const BBox& box2) {
|
||
float x1 = std::max(box1.x1, box2.x1);
|
||
float y1 = std::max(box1.y1, box2.y1);
|
||
float x2 = std::min(box1.x2, box2.x2);
|
||
float y2 = std::min(box1.y2, box2.y2);
|
||
|
||
float intersection = std::max(0.0f, x2 - x1) * std::max(0.0f, y2 - y1);
|
||
float area1 = (box1.x2 - box1.x1) * (box1.y2 - box1.y1);
|
||
float area2 = (box2.x2 - box2.x1) * (box2.y2 - box2.y1);
|
||
|
||
return intersection / (area1 + area2 - intersection);
|
||
}
|
||
|
||
bool TrtInference::createContext() {
|
||
if (!engine_) {
|
||
return false;
|
||
}
|
||
|
||
context_ = engine_->createExecutionContext();
|
||
if (!context_) {
|
||
return false;
|
||
}
|
||
|
||
// 设置优化的推理维度
|
||
if (config_.max_batch_size > 0) {
|
||
// TensorRT 8.x 不再使用 setOptimizationProfile
|
||
context_->setOptimizationProfileAsync(0, stream_->get());
|
||
}
|
||
|
||
return true;
|
||
}
|
||
|
||
bool TrtInference::allocateBuffers() {
|
||
if (!engine_ || !context_) {
|
||
return false;
|
||
}
|
||
|
||
try {
|
||
// 先释放旧的缓冲区
|
||
input_buffers_.clear();
|
||
output_buffers_.clear();
|
||
bindings_.clear();
|
||
output_names_.clear();
|
||
|
||
// 获取网络的输入输出数量
|
||
int numInputs = 0;
|
||
int numOutputs = 0;
|
||
int maxBindings = engine_->getNbIOTensors();
|
||
bindings_.resize(maxBindings);
|
||
|
||
// 分配输入输出缓冲区
|
||
for (int i = 0; i < maxBindings; i++) {
|
||
const char* name = engine_->getIOTensorName(i);
|
||
nvinfer1::TensorIOMode mode = engine_->getTensorIOMode(name);
|
||
if (mode == nvinfer1::TensorIOMode::kINPUT) {
|
||
numInputs++;
|
||
nvinfer1::Dims dims = engine_->getTensorShape(name);
|
||
nvinfer1::DataType dtype = engine_->getTensorDataType(name);
|
||
|
||
// 计算缓冲区大小
|
||
size_t size = 1;
|
||
for (int j = 0; j < dims.nbDims; j++) {
|
||
if (dims.d[j] < 0) {
|
||
dims.d[j] = (j == 0) ? 1 : config_.model.input_shape[j-1];
|
||
}
|
||
size *= dims.d[j];
|
||
}
|
||
size *= (dtype == nvinfer1::DataType::kFLOAT ? 4 : 2); // FP32或FP16
|
||
|
||
// 使用固定内存分配
|
||
void* d_ptr = nullptr;
|
||
cudaMalloc(&d_ptr, size);
|
||
if (!d_ptr) {
|
||
Logger::error("Failed to allocate input buffer");
|
||
return false;
|
||
}
|
||
|
||
input_buffers_.emplace_back(std::make_unique<detail::CudaBuffer>(size));
|
||
bindings_[i] = d_ptr;
|
||
|
||
} else if (mode == nvinfer1::TensorIOMode::kOUTPUT) {
|
||
numOutputs++;
|
||
nvinfer1::Dims dims = engine_->getTensorShape(name);
|
||
nvinfer1::DataType dtype = engine_->getTensorDataType(name);
|
||
|
||
// 计算缓冲区大小
|
||
size_t size = 1;
|
||
for (int j = 0; j < dims.nbDims; j++) {
|
||
if (dims.d[j] < 0) {
|
||
dims.d[j] = (j == 0) ? 1 : config_.model.input_shape[j-1];
|
||
}
|
||
size *= dims.d[j];
|
||
}
|
||
size *= (dtype == nvinfer1::DataType::kFLOAT ? 4 : 2); // FP32或FP16
|
||
|
||
// 使用固定内存分配
|
||
void* d_ptr = nullptr;
|
||
cudaMalloc(&d_ptr, size);
|
||
if (!d_ptr) {
|
||
Logger::error("Failed to allocate output buffer");
|
||
return false;
|
||
}
|
||
|
||
output_buffers_.emplace_back(std::make_unique<detail::CudaBuffer>(size));
|
||
bindings_[i] = d_ptr;
|
||
output_names_.push_back(name);
|
||
}
|
||
}
|
||
|
||
// 同步设备
|
||
cudaDeviceSynchronize();
|
||
|
||
return (numInputs > 0 && numOutputs > 0);
|
||
}
|
||
catch (const std::exception& e) {
|
||
Logger::error("Error in allocateBuffers: " + std::string(e.what()));
|
||
return false;
|
||
}
|
||
}
|
||
|
||
void TrtInference::destroy() {
|
||
try {
|
||
// 先同步CUDA流和设备
|
||
if (stream_) {
|
||
stream_->sync();
|
||
}
|
||
cudaDeviceSynchronize();
|
||
|
||
// 释放绑定的设备内存
|
||
for (void* ptr : bindings_) {
|
||
if (ptr) {
|
||
cudaFree(ptr);
|
||
}
|
||
}
|
||
|
||
// 释放缓冲区
|
||
input_buffers_.clear();
|
||
output_buffers_.clear();
|
||
bindings_.clear();
|
||
output_names_.clear();
|
||
|
||
// 释放TensorRT资源
|
||
if (context_) {
|
||
delete context_; // 使用delete而不是destroy()
|
||
context_ = nullptr;
|
||
}
|
||
if (engine_) {
|
||
delete engine_; // 使用delete而不是destroy()
|
||
engine_ = nullptr;
|
||
}
|
||
if (runtime_) {
|
||
delete runtime_; // 使用delete而不是destroy()
|
||
runtime_ = nullptr;
|
||
}
|
||
|
||
// 最后再次同步设备
|
||
cudaDeviceSynchronize();
|
||
}
|
||
catch (const std::exception& e) {
|
||
Logger::error("Error in destroy: " + std::string(e.what()));
|
||
}
|
||
}
|
||
|
||
} // namespace pipeline
|