【实操】以制作通用型TensorRT推理二进制(DLL/SO)文件为例讲解接口封装

80 阅读5分钟

关键词: yolov8-pose、DLL|SO、TENSORRT、ENGINE、TRT、C++、PYTHON

前言

  最近整理关于GPU端的TensorRT的加速框架时倍感繁琐,感叹是否可以制作一个二进制文件便于后期直接推理调用。由此便诞生了今天的博客文章记录这次制作过程中所踩的坑以及记录解决方案。

TensorRT推理封装分析

  在这里我们以YOLOV8-POSE为例子进行讲解整体的TRT推理。整体的推理步骤大致可分为:

  1. 读取权重、配置iou 、conf 、point等基础参数
  2. 权重预加热部分;
  3. 图像进行预处理部分;
  4. 图像数据推理部分;
  5. 图像数据后处理解码部分;
  6. 解码数据进行绘制到图像部分;

对于大部分的推理框架来说,例如pytorch、onnxruntime和OpenVION框架而言推理大致步骤亦是如此。通过分析上述步骤我们可以将整个推理步骤进行封装为如下几步:

  1. 初始化步骤:该步骤主要为设定权重路径、模型尺寸、各种阈值设定以及权重的预加热步骤;
  2. 推理步骤:该步骤主要是接收输入的图像数据并将初始化步骤的阈值进行计算得到bboxes和points数据;
  3. 关闭步骤:该步骤主要是对加载的一些数据和创建的内存区域进行置空释放,避免占用内存导致内存溢出或内存泄漏等原因

封装步骤实现

  关于整个yolov8-pose的推理步骤在github中有很多,大家可自行下载并调试运行,这里我将着重讲解对于封装过程的一些实现步骤以及踩坑的地方以及解决办法。这里我们以win系统下的生成dll为例子进行讲解(细节上win和linux库的差异)初始化步骤的实现,

初始化步骤

  这一步较为简单,需要注意的地方是全局变量的声明,这里的YOLOv8_pose* V8Pose 、cv::Size size在多个函数中会用到,放在全局变量中进行定义会比较好。

#include "opencv2/opencv.hpp"
#include "yolov8-pose.hpp"
#include <chrono>

// 全局变量
YOLOv8_pose* V8Pose = nullptr;
cv::Mat res;
cv::Size size;
std::vector<Object> objs;

// 初始化 YOLOv8_pose 和 图像
extern "C" __declspec(dllexport) void init_v8pose(const char* engine_file_path, int model_width, int model_height) {
    std::string engine_path(engine_file_path);
    cudaSetDevice(0);
    V8Pose = new YOLOv8_pose(engine_path); // 加载pose权重
    V8Pose->make_pipe(true); // 预处理
    size = cv::Size{model_width, model_height};
}

推理步骤

  最重要推理步骤是这篇博客踩坑较多的地方。由于我们在初始化步骤中明确了使用CUDA(以及使用是哪块)整个v8pose声明了,权重也进行调用了,那么这步就需要对输入的图像当作输入另外加上计算设置的阈值部分进行当作输入。但是后续的分析实现的过程中我发现返回值并不好构建,这里的返回值大多是由如下三部分构成:

  1. 绘制好预测结果的图像;
  2. 计算得到目标的检测框坐标集;
  3. 计算得到关键点坐标集;

  分析上述的三步骤我们有了三个返回值,其中一个是mat格式数据,另外两个是std::vector格式的数据,这里首先相当的则是构建一个结构体,在绘制步骤即可获取上述的三个输出结果。 结构体的构建:

struct DrawOut {
    std::vector<int> bboxe; // bboxe数据
    std::vector<int> point; // point数据
    int bboxe_size;         // bboxe数据的大小
    int point_size;         // point数据的大小

};

  上述的结构体当作返回值我们很好就能处理在C++端进行调用二进制文件,但是这未免对python的同事有些苛刻,他们不便于对结构体和vector数据进行调用。那么将vector转为数组类型对python端和C\C++端都能满足二者的需求。

  通过编写C++示例生产数组给python进行调用,此时我们发现,对返回值进行读取的时候存在一个纰漏:数组的大小位置。由于C语言中对于数组的构建是较为严谨的,这和python中的数组构建存在差异,因此在C\C++端必须将数组的大小一并返回。 这里我们又不能使用结构体进行当作返回值进行操作,观察函数的构建我们在输入端是可以分开构建输入进行调用,岂不是可以在函数的构建中对返回信息取地址赋予输入值中解决了实际问题么!

// 处理图像并获取绘制信息
extern "C" __declspec(dllexport) void processImage(float score_thres, 
    float iou_thres, 
    float point_thres,
    int image_width,
    int image_height,
    unsigned char* image_data,
    unsigned char** res_data, 
    int** bboxes, 
    int**  points, 
    int* bboxes_size, 
    int* points_size) 
{
    int topk = 100;
    cv::Mat image(image_width, image_height, CV_8UC3, image_data, 0);//为了显示图片 先改为mat类型  三维图
    V8Pose->copy_from_Mat(image, size);
    V8Pose->infer();
    V8Pose->postprocess(objs, score_thres, iou_thres, topk);
    DrawOut my_out = V8Pose->draw_objects(image, res, objs, point_thres);
    *res_data = res.data;
    *bboxes = my_out.bboxe.data();
    *points = my_out.point.data();
    *bboxes_size = my_out.bboxe_size;
    *points_size = my_out.point_size;
}

关闭步骤

  经过调试上述操作中需要对其内存进行释放可得到如下:

// 清空中间缓存
extern "C" __declspec(dllexport) void clear_process() {
    if (!res.empty()) {
        res.release();
    }
    size = cv::Size();
    objs.clear();
    cv::destroyAllWindows();
}

// 清理资源
extern "C" __declspec(dllexport) void clear_v8pose() {
    delete V8Pose;
    V8Pose = nullptr;
}

总结

  该二进制文件的制作思路献于诸位,愿诸公工作顺利。