赞
踩
1.首先是生成tensorRT的engine文件:
Tensorrt的环境配置并生成tensorrt文件参考与以下的B站视频链接。
【手把手带你实战YOLOv5-部署篇】YOLOv5 TensorRT环境安装与配置_哔哩哔哩_bilibili
其中遇到了不少的问题,记录如下:
(1)AssertionError: Invalid CUDA '--device 0' requested, use '--device cpu' or pass valid CUDA device(s)
这个问题是说cuda不可用,最终发现就是cuda没有正确安装,要安装cuda的GPU版本, cudnn,并且torch和torchvision的版本都要与cuda对应,缺一不可。
(2)tensorrt: export failure 0.0s: No module named 'tensorrt':问题是说tensorrt没有正确安装成功,但是我是根据视频的教程一步一步来的,于是我就想卸载重装,但是在终端中删除这个库的时候出现了错误,说我没装这个库,警告的信息如下。
WARNING: Ignoring invalid distribution -pencv-python (d:\anaconda3\lib\site-packages)
WARNING: Skipping tensorrt as it is not installed.
与是我发现这个是我base环境的位置,于是我换了个思路,在base环境中也用pip install的指令安装了对应版本的tensorrt发现能够正常运行了。
以上终于能够实现.pt模型转.engine文件。
2.在C++项目中调用.engine文件
首先要进行tensorRT、OpenCV的环境配置,这个找不到参考的文献了,就是配置一下项目的库目录、包含目录、以及附加依赖项。
在配置好环境之后,读取文件模型:
- // 读取本地模型文件
- std::string model_path_engine = model_path;
- std::ifstream file_ptr(model_path_engine, std::ios::binary);
- if (!file_ptr.good()) {
- std::cerr << "文件无法打开,请确定文件是否可用!" << std::endl;
- }
- size_t size = 0;
- file_ptr.seekg(0, file_ptr.end);
- size = file_ptr.tellg();
- file_ptr.seekg(0, file_ptr.beg);
- char* model_stream = new char[size];
- file_ptr.read(model_stream, size);
- file_ptr.close();
做一些推理前的准备工作,并分配内存
- // 反序列化引擎
- nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
- // 推理引擎 // 保存模型的模型结构、模型参数以及最优计算kernel配置; // 不能跨平台和跨TensorRT版本移植
- nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(model_stream, size);
- // 上下文// 储存中间值,实际进行推理的对象// 由engine创建,可创建多个对象,进行多推理任务
- nvinfer1::IExecutionContext* context = engine->createExecutionContext();
- delete[] model_stream;
- // 创建GPU显存缓冲区
- void** data_buffer = new void* [num_ionode];
- // 创建GPU显存输入缓冲区
- int input_node_index = engine->getBindingIndex(input_node_name);
- nvinfer1::Dims input_node_dim = engine->getBindingDimensions(input_node_index);
- size_t input_data_length = input_node_dim.d[1] * input_node_dim.d[2] * input_node_dim.d[3];
- cudaMalloc(&(data_buffer[input_node_index]), input_data_length * sizeof(float));
- // 创建GPU显存输出缓冲区
- int output_node_index = engine->getBindingIndex(output_node_name);
- nvinfer1::Dims output_node_dim = engine->getBindingDimensions(output_node_index);
- size_t output_data_length = output_node_dim.d[1] * output_node_dim.d[2];
- cudaMalloc(&(data_buffer[output_node_index]), output_data_length * sizeof(float));
创建cuda流:
- // 创建输入cuda流
- cudaStream_t stream;
- cudaStreamCreate(&stream);
- std::vector<float> input_data(input_data_length);
- memcpy(input_data.data(), BN_image.ptr<float>(), input_data_length * sizeof(float));
模型推理:
- // 模型推理、记录模型推理时间
- auto start_time = std::chrono::high_resolution_clock::now();
- context->enqueueV2(data_buffer, stream, nullptr); //这里得到的data_buffer就是模型的输入和输出
- auto end_time = std::chrono::high_resolution_clock::now();
- std::chrono::duration<double, std::milli> inference_time = end_time - start_time;
- std::cout << "-----------------Inference Time: " << inference_time.count() << " ms------------------" << std::endl;
- float* result_array = new float[output_data_length];
- // 输出数据由GPU转到CPU
- cudaMemcpyAsync(result_array, data_buffer[output_node_index], output_data_length * sizeof(float), cudaMemcpyDeviceToHost, stream);
之后利用模型的输出进行框的解码工作就行了。解码的代码是参考一个大佬的,网址我也找不到了,很抱歉。框的解码包括了输出解码、置信度过滤、非极大值抑制、最后绘制检测框的内容。
- cv::Mat net_result(cv::Mat image, float* result, float factor, std::vector<std::string> class_names)
- {
- cv::Mat det_output = cv::Mat(25200, 85, CV_32F, result);
- std::vector<cv::Rect> position_boxes;
- std::vector<int> classIds;
- std::vector<float> confidences;
-
- std::cout << det_output.rows << std::endl; //共25200个框,每个框有85个信息
- for (int i = 0; i < det_output.rows; i++) {
- float confidence = det_output.at<float>(i, 4);
- if (confidence < 0.2) {
- continue;
- }
- std::cout << "confidence: " << confidence << std::endl;
- cv::Mat classes_scores = det_output.row(i).colRange(5, 85);
- cv::Point classIdPoint;
- double score;
- // 获取一组数据中最大值及其位置
- minMaxLoc(classes_scores, 0, &score, 0, &classIdPoint);
- // 置信度 0~1之间
- if (score > 0.25)
- {
- float cx = det_output.at<float>(i, 0);
- float cy = det_output.at<float>(i, 1);
- float ow = det_output.at<float>(i, 2);
- float oh = det_output.at<float>(i, 3);
- int x = static_cast<int>((cx - 0.5 * ow) * factor);
- int y = static_cast<int>((cy - 0.5 * oh) * factor);
- int width = static_cast<int>(ow * factor);
- int height = static_cast<int>(oh * factor);
- cv::Rect box;
- box.x = x;
- box.y = y;
- box.width = width;
- box.height = height;
-
- position_boxes.push_back(box);
- classIds.push_back(classIdPoint.x);
- confidences.push_back(score);
- }
- }
- // NMS
- std::vector<int> indexes;
- cv::dnn::NMSBoxes(position_boxes, confidences, 0.25, 0.45, indexes);
- for (size_t i = 0; i < indexes.size(); i++) {
- int index = indexes[i];
- int idx = classIds[index];
- cv::rectangle(image, position_boxes[index], cv::Scalar(0, 0, 255), 2, 8);
- cv::rectangle(image, cv::Point(position_boxes[index].tl().x, position_boxes[index].tl().y - 20),
- cv::Point(position_boxes[index].br().x, position_boxes[index].tl().y), cv::Scalar(0, 255, 255), -1);
- cv::putText(image, class_names[idx], cv::Point(position_boxes[index].tl().x, position_boxes[index].tl().y - 10), cv::FONT_HERSHEY_SIMPLEX, .5, cv::Scalar(0, 0, 0));
- }
- return image;
- }
写得很烂,主要是自己想记录一下,后续如果找到了参考文献的话会贴在评论区里面,大佬写得比我好多了。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。