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;
}
写得很烂,主要是自己想记录一下,后续如果找到了参考文献的话会贴在评论区里面,大佬写得比我好多了。