赞
踩
最近开始机器视觉、深度学习在Linux c++方向(大概就是这个方向)的学习,其实主要是为了学习c++在Linux环境做一些工程,刚好自己是做机器视觉相关算法的,但之前大部分用的都是python,以及jetson盒子相关配套的东西,所以就从TensorRT-C++推理开始学习了。
这里也是记录一下自己的学习过程,因为c++较为繁琐复杂,通过笔记来记录加深印象。
首先贴出学习代码的来源:
TensorRT-For-YOLO-Series 后来发现该代码应该是参照YoloX写的
YoloX-TRT
一些环境的安装,譬如我是用vscode在linux远程服务器建立docker 容器编写、编译、运行c++代码的,这些也都参考了很多回答和文章,如果有需要可以评论我专门写一篇文章用于介绍前面的准备工作。
一切都从环境安装开始,所幸本次学习的代码环境安装并不复杂。
机器配置就是普通的3090两卡服务器,采用的是docker镜像做容器的方式配置的环境。
docker镜像: nvcr.io/nvidia/tensorrt:23.10-py3 (docker镜像版本里面的cuda版本要能与宿主机的cuda driver适应)
起容器命令: sudo docker run -it -v /data/smh:/workdir --workdir=/workdir/ --network=host --runtime nvidia --gpus all --shm-size 16g --name trt nvcr.io/nvidia/tensorrt:23.10-py3 /bin/bash
进去之后看一下cmake版本,cmake --version,我这里是3.24.0。
OpenCV安装 :可以选择自己安装编译opencv的包(这个需要自己多查查,我也是查了好久结合好几个回答,才安装编译成功),但是这样的话需要修改工程里面cmakelist相关路径,推荐不太懂cmake的伙伴直接通过系统安装。
可能需要安装的依赖(g++,python库,读取视频流的库等,可以先试试安装opencv的命令,出问题再回来执行安装这些库):
apt update && apt upgrade
apt-get install build-essential libgtk2.0-dev libgtk-3-dev libavcodec-dev libavformat-dev libjpeg-dev libswscale-dev libtiff5-dev
apt install python3-dev python3-numpy
apt install libgstreamer-plugins-base1.0-dev libgstreamer1.0-dev
apt install libpng-dev libopenexr-dev libtiff-dev libwebp-dev
系统安装OpenCV的命令:apt-get update -y && apt-get install -y libopencv-dev
至此,TensorRT相关的包,拉取英伟达的镜像里面基本都安装好了,大概率需要自己安装opencv,其余的应该就没什么了,所以环境搭建基本完成。
TensorRT-For-YOLO-Series/cpp/end2end/main.cpp
首先是类的一些基本定义,这里一共两个类 Logger和yolo。
Logger类继承自nvinfer1::ILogger,通过实现ILogger接口,并将其传递给TensorRT相关的API,可以捕获和处理TensorRT运行时产生的日志消息。
这个ILogger类是TensorRT建立、使用各种api的基础。
// Logger类,继承自父类 nvinfer1::ILogger
class Logger : public ILogger {
public:
// noexcept override 表示不允许被重载
void log(Severity severity, const char* msg) noexcept override {
if (severity != Severity::kINFO) {
std::cout << msg << std::endl;
}
}
};
yolo类是针对于yolo模型定义的类,里面主要包含一些对于图片的预处理,加载trt模型、推理、销毁内存等。
class Yolo { public: Yolo(char* model_path); // 构造函数,初始化空间、cuda流等等 // letterbox用于resize图片并填充图片 float letterbox( const cv::Mat& image, cv::Mat& out_image, const cv::Size& new_shape, int stride, const cv::Scalar& color, bool fixed_shape, bool scale_up); float* blobFromImage(cv::Mat& img); //相当于标准化,并且将图片转化为一个float数组 void draw_objects(const cv::Mat& img, float* Boxes, int* ClassIndexs, int* BboxNum); void Init(char* model_path); void Infer( int aWidth, int aHeight, int aChannel, unsigned char* aBytes, float* Boxes, int* ClassIndexs, int* BboxNum); ~Yolo(); private: nvinfer1::ICudaEngine* engine = nullptr; nvinfer1::IRuntime* runtime = nullptr; nvinfer1::IExecutionContext* context = nullptr; cudaStream_t stream = nullptr; void* buffs[5]; // 为输入输出开辟的内存空间 int iH, iW, in_size, out_size1, out_size2, out_size3, out_size4; Logger gLogger; // 要创建Builder、Runtime等类的实例,必须先初始化ILogger实例 };
赋值构造函数只需要一个参数:模型地址,加载模型并反序列化到buffer,再初始化一些cuda流等等。
Yolo::Yolo(char* model_path) { ifstream ifile(model_path, ios::in | ios::binary); // 读取的文件流 if (!ifile) { cout << "read serialized file failed\n"; std::abort(); // 直接发送异常,终止程序 } ifile.seekg(0, ios::end); // seekg设置指针位置,这里设置到文件末尾 const int mdsize = ifile.tellg(); // tellg是当前指针位置,这里计算出文件的size // 多次操作同一个文件流,为了避免后面的操作因为文件流处于eof状态而无法正常进行,在每种操作方法前,都使用clear()方法将文件流状态重设为默认的“goodbit”状态。 ifile.clear(); ifile.seekg(0, ios::beg); vector<char> buf(mdsize); ifile.read(&buf[0], mdsize); // 把文件流读取到buf ifile.close(); cout << "model size: " << mdsize << endl; runtime = nvinfer1::createInferRuntime(gLogger); //创建一个Runtime实例 initLibNvInferPlugins(&gLogger, ""); //调用initLibNvInferPlugins 注册所有插件 engine = runtime->deserializeCudaEngine((void*)&buf[0], mdsize); // 反序列化模型,入参是加载模型的buffer和模型的size // getBindingIndex通过tensor的name获取tensor的序号,再传入getBindingDimensions获取该tensor的shape auto in_dims = engine->getBindingDimensions(engine->getBindingIndex("images")); // auto dd = engine->getTensorShape("images"); iH = in_dims.d[2]; iW = in_dims.d[3]; // cout << dd.d[2] << in_dims.d[2] << dd.d[3] << in_dims.d[3] << endl; // 计算输入输出tensor需要申请的空间size in_size = 1; for (int j = 0; j < in_dims.nbDims; j++) { in_size *= in_dims.d[j]; } auto out_dims1 = engine->getBindingDimensions(engine->getBindingIndex("num")); out_size1 = 1; for (int j = 0; j < out_dims1.nbDims; j++) { out_size1 *= out_dims1.d[j]; } auto out_dims2 = engine->getBindingDimensions(engine->getBindingIndex("boxes")); out_size2 = 1; for (int j = 0; j < out_dims2.nbDims; j++) { out_size2 *= out_dims2.d[j]; } auto out_dims3 = engine->getBindingDimensions(engine->getBindingIndex("scores")); out_size3 = 1; for (int j = 0; j < out_dims3.nbDims; j++) { out_size3 *= out_dims3.d[j]; } auto out_dims4 = engine->getBindingDimensions(engine->getBindingIndex("classes")); out_size4 = 1; for (int j = 0; j < out_dims4.nbDims; j++) { out_size4 *= out_dims4.d[j]; } context = engine->createExecutionContext(); // 创建一个执行上下文(Execution Context) if (!context) { cout << "create execution context failed\n"; std::abort(); } /* 申请空间 cudaMalloc根据指定的大小,从 GPU 的全局内存中分配一块连续的内存空间,并返回一个指向该内存空间的指针。 这样就可以在 GPU 上使用该指针来访问和操作所分配的内存空间。 使用 cudaMalloc() 方法分配的内存需要在使用完毕后通过 "cudaFree()" 方法进行释放,以避免内存泄漏。 */ cudaError_t state; state = cudaMalloc(&buffs[0], in_size * sizeof(float)); if (state) { cout << "allocate memory failed\n"; std::abort(); } state = cudaMalloc(&buffs[1], out_size1 * sizeof(int)); if (state) { cout << "allocate memory failed\n"; std::abort(); } state = cudaMalloc(&buffs[2], out_size2 * sizeof(float)); if (state) { cout << "allocate memory failed\n"; std::abort(); } state = cudaMalloc(&buffs[3], out_size3 * sizeof(float)); if (state) { cout << "allocate memory failed\n"; std::abort(); } state = cudaMalloc(&buffs[4], out_size4 * sizeof(int)); if (state) { cout << "allocate memory failed\n"; std::abort(); } // 通过 cudaStreamCreate() 方法,可以创建一个 CUDA 流对象,并返回一个对应的流句柄(stream handle)。 // CUDA 流可以用于在 GPU 上并行执行多个操作,例如内存传输、内核函数调用等。通过将不同的操作分配到不同的 CUDA 流中,可以实现这些操作的并发执行,从而提高程序的性能。 state = cudaStreamCreate(&stream); if (state) { cout << "create stream failed\n"; std::abort(); } }
letterbox函数用于resize输入图片到模型输入的尺寸,并以灰边填充。这样做的好处是保持原本图片的比例,暴力resize可能会导致图片失真等等。
可以保存中间图片看一下效果:
letterbox的思路和做法可以总结为:
float Yolo::letterbox( const cv::Mat& image, // 输入图片 cv::Mat& out_image, // resize填充之后的输出图片 const cv::Size& new_shape = cv::Size(640, 640), // resize的size int stride = 32, const cv::Scalar& color = cv::Scalar(114, 114, 114), // 填充的颜色 bool fixed_shape = false, bool scale_up = true) { // scale_up 参数用来控制图片是否增大 cv::Size shape = image.size(); // 先获取输入图片原本的size // cout << "image height: " << shape.height << ", image width: " << shape.width << endl; // 计算一下图片原本size和目标size的比例,取小的那个 float r = std::min( (float)new_shape.height / (float)shape.height, (float)new_shape.width / (float)shape.width); // scale_up 参数用来控制图片是否增大 // 如果图片本身比设置的new_size小,需要通过scale_up参数来判断是否resize(放大),或保持不变 if (!scale_up) { r = std::min(r, 1.0f); } //cout << "r: " << r << endl; // 如果输入图像h:1440,w:2560,resize后会变成h:360,w:640 int newUnpad[2]{ (int)std::round((float)shape.width * r), (int)std::round((float)shape.height * r)}; // 创建一个size为 new shape的 Mat tmp -> resize原图得到的 cv::Mat tmp; if (shape.width != newUnpad[0] || shape.height != newUnpad[1]) { cv::resize(image, tmp, cv::Size(newUnpad[0], newUnpad[1])); } else { tmp = image.clone(); } // 先按最小的比例resize, float dw = new_shape.width - newUnpad[0]; float dh = new_shape.height - newUnpad[1]; // cout << "newUnpad[0]: " << newUnpad[0] << ", newUnpad[1]: " << newUnpad[1] << endl; // cout << "dh1: " << dh << ", dw1: " << dw << endl; // fixed_shape 不清楚功能,可能是让填充之后的size能够是32的倍数 if (!fixed_shape) { dw = (float)((int)dw % stride); dh = (float)((int)dh % stride); // cout << "dh2: " << dh << ", dw2: " << dw << endl; } // copyMakeBorder接收参数是向四周扩展的像素数,所以宽高分别除以2 dw /= 2.0f; dh /= 2.0f; // cout << "dh3: " << dh << ", dw3: " << dw << endl; int top = int(std::round(dh - 0.1f)); int bottom = int(std::round(dh + 0.1f)); int left = int(std::round(dw - 0.1f)); int right = int(std::round(dw + 0.1f)); cv::copyMakeBorder(tmp, out_image, top, bottom, left, right, cv::BORDER_CONSTANT, color); cv::imwrite("letterbox1.jpg", tmp); cv::imwrite("letterbox2.jpg", out_image); return 1.0f / r; }
blobFromImage函数就是将img拉长放在一个float数组里,且标准化。draw_objects函数将检测结果绘制在图片上,用到的也都是一些opencv的函数。
Infer函数是模型执行推理的函数,主要就是遵循TRT和CUDA那一套流程:
void Yolo::Infer( int aWidth, //图片的宽 int aHeight, //图片的高 int aChannel, //图片的通道数 unsigned char* aBytes, //图片的数据 cv::Mat data值 float* Boxes, // 用于存放结果的数组指针 int* ClassIndexs, int* BboxNum) { cv::Mat img(aHeight, aWidth, CV_MAKETYPE(CV_8U, aChannel), aBytes); // 输入的图片 cv::Mat pr_img; // 转换后的图片先预定义 float scale = letterbox(img, pr_img, {iW, iH}, 32, {114, 114, 114}, true); // 通过letterbox进行转换 cv::cvtColor(pr_img, pr_img, cv::COLOR_BGR2RGB); // 图片BGR转RGB float* blob = blobFromImage(pr_img); //执行blob // 静态变量 static int* num_dets = new int[out_size1]; // 对应检测到的每个目标 static float* det_boxes = new float[out_size2]; // 每个目标对应的box static float* det_scores = new float[out_size3]; // 每个目标的置信度得分 static int* det_classes = new int[out_size4]; // 每个目标的类别 // cudaMemcpyAsync 是 CUDA 库中的一个函数,用于在 GPU 之间或在主机和设备之间异步地进行内存数据的传输。 // 使用 cudaMemcpyAsync 函数进行内存数据传输时,数据传输是异步执行的,即函数调用会立即返回,而不会等待数据传输完成。 // 因此,可以在数据传输进行的同时进行其他的 GPU 计算操作,以提高程序的效率。 // 这里相当于把blob上的数据复制到buffs上 cudaError_t state = cudaMemcpyAsync(buffs[0], &blob[0], in_size * sizeof(float), cudaMemcpyHostToDevice, stream); if (state) { cout << "transmit to device failed\n"; std::abort(); } // enqueueV2 是 nvinfer1::IExecutionContext::enqueueV2方法,有三个参数 // bindings 输入或输出array的缓存的指针 // stream 一个cuda stream // inputConsumed 输入的buffer可以被重新填充的信号(不清楚怎么使用) // 个人理解就是执行推理 context->enqueueV2(&buffs[0], stream, nullptr); // 推理结束之后,把device上的output数据再传回host state = cudaMemcpyAsync(num_dets, buffs[1], out_size1 * sizeof(int), cudaMemcpyDeviceToHost, stream); if (state) { cout << "transmit to host failed \n"; std::abort(); } state = cudaMemcpyAsync( det_boxes, buffs[2], out_size2 * sizeof(float), cudaMemcpyDeviceToHost, stream); if (state) { cout << "transmit to host failed \n"; std::abort(); } state = cudaMemcpyAsync( det_scores, buffs[3], out_size3 * sizeof(float), cudaMemcpyDeviceToHost, stream); if (state) { cout << "transmit to host failed \n"; std::abort(); } state = cudaMemcpyAsync( det_classes, buffs[4], out_size4 * sizeof(int), cudaMemcpyDeviceToHost, stream); if (state) { cout << "transmit to host failed \n"; std::abort(); } BboxNum[0] = num_dets[0]; int img_w = img.cols; int img_h = img.rows; // cout << "iW: " <<iW << "iH: " <<iH << "img_w: " <<img_w << "img_h: " <<img_h << "scale: " <<scale << endl; // 计算一下宽高不同的图片转换回来的偏差,因为上面的letterbox是按最小的scale缩放的 int x_offset = (iW * scale - img_w) / 2; int y_offset = (iH * scale - img_h) / 2; for (size_t i = 0; i < num_dets[0]; i++) { float x0 = (det_boxes[i * 4]) * scale - x_offset; float y0 = (det_boxes[i * 4 + 1]) * scale - y_offset; float x1 = (det_boxes[i * 4 + 2]) * scale - x_offset; float y1 = (det_boxes[i * 4 + 3]) * scale - y_offset; // 防止框的点越界 x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f); y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f); x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f); y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f); Boxes[i * 4] = x0; Boxes[i * 4 + 1] = y0; Boxes[i * 4 + 2] = x1 - x0; Boxes[i * 4 + 3] = y1 - y0; ClassIndexs[i] = det_classes[i]; } delete blob; }
析构函数就是一些内存、实例的销毁
Yolo::~Yolo() {
cudaStreamSynchronize(stream); // 用于同步等待指定的 CUDA 流上的所有操作完成。
cudaFree(buffs[0]);
cudaFree(buffs[1]);
cudaFree(buffs[2]);
cudaFree(buffs[3]);
cudaFree(buffs[4]);
cudaStreamDestroy(stream);
context->destroy();
engine->destroy();
runtime->destroy();
}
最后就是主程序调用实验了,有两个入参:trt模型路径和需要推理的图片路径
int main(int argc, char** argv) { // 输入参数有-model_path 和 -image_path if (argc == 5 && std::string(argv[1]) == "-model_path" && std::string(argv[3]) == "-image_path") { char* model_path = argv[2]; char* image_path = argv[4]; float* Boxes = new float[4000]; int* BboxNum = new int[1]; int* ClassIndexs = new int[1000]; Yolo yolo(model_path); cv::Mat img; img = cv::imread(image_path); // warmup for (int num =0; num < 10; num++) { yolo.Infer(img.cols, img.rows, img.channels(), img.data, Boxes, ClassIndexs, BboxNum); } // run inference auto start = std::chrono::system_clock::now(); yolo.Infer(img.cols, img.rows, img.channels(), img.data, Boxes, ClassIndexs, BboxNum); auto end = std::chrono::system_clock::now(); std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl; yolo.draw_objects(img, Boxes, ClassIndexs, BboxNum); } else { std::cerr << "--> arguments not right!" << std::endl; std::cerr << "--> yolo -model_path ./output.trt -image_path ./demo.jpg" << std::endl; return -1; } }
回忆了很多c++的基本知识,也对齐了python用TRT推理和c++用TRT的一个流程上的异同,并且在查找很多TRT的c++ api时发现很多api弃用了或者推荐使用新的api,后续计划将该工程对应的api修改换成新版的作为一个课后补习。
如有错误,请联系我修改。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。