CGO调用TensorRT

shelgi
200 阅读8分钟

1. 前言

之前写过一篇Go调用C++动态库实现车牌识别，当时是自己业余写着玩的demo，代码中存在不少问题。后续呢自己在压测过程中也发现了内存泄漏等Bug。所以这一篇Blog就来对之前的一些错误进行更正，同时完善一下整个流程。

2. CGO部分

在这部分主要是讲讲c++中的头文件定义以及go对动态链接库的调用，目前呢我是有两种不同的方式来传递推理图像

  • 使用gocv读取图片，然后向c++部分传递Mat指针
  • 使用base64传递，在c++部分进行base64转码为Mat

由于上一版的实现是基于第二种方式，所以这一节主要就讲第一种方式实现。另外，由于本文主要是想说明cgo的使用，所以关于TensorRT推理部分的代码我打算单独写一篇，避免内容混杂。

2.1 C++ 头文件定义及实现

#pragma once
#include<stdio.h>
#include <stdlib.h>
​
#ifdef __cplusplus
extern "C"{
#endif
    void* yolo(const char* enginePath,const char* labelPath,float score_thresh=0.25f,float iou_thresh=0.4f,int deviceId=0);
    void release(void* model);
    char* detect(void* model,void* image,int* count);
#ifdef __cplusplus
}
#endif

这个便是头文件wrap.h的定义，主要就是构建engine以及定义推理的函数。这里因为我把前后处理全部封装到了cuda算子中，所以在构建engine的时候就需要传入置信度以及iou阈值去定义nms的后处理插件部分。对应的yolo.h就是稍微封装一下engine以及对应labels，这里我拿一个obb模型为例。

#pragma once
#include <vector>
#include <string>
#include "yolov8_obb_detector.h"
​
class Yolo {
public:
    Yolo() = default;
    Yolo(YoloOBB::Infer* engine,std::vector<std::string> labels):engine(engine),labels(std::move(labels)){}
    ~Yolo() {
        if(engine) {
            delete engine;
            engine = nullptr;
        }
    };
    YoloOBB::Infer* engine = nullptr;
    std::vector<std::string> labels;
};

最后再来看看wrap.cpp中的具体实现

static std::vector<cv::Point> xywhr2xyxyxyxy(const YoloOBB::Box& box)
{
    float cos_value = std::cos(box.angle);
    float sin_value = std::sin(box.angle);
​
    float w_2 = box.width / 2, h_2 = box.height / 2;
    float vec1_x =  w_2 * cos_value, vec1_y = w_2 * sin_value;
    float vec2_x = -h_2 * sin_value, vec2_y = h_2 * cos_value;
​
    std::vector<cv::Point> corners;
    corners.emplace_back(box.center_x + vec1_x + vec2_x, box.center_y + vec1_y + vec2_y);
    corners.emplace_back(box.center_x + vec1_x - vec2_x, box.center_y + vec1_y - vec2_y);
    corners.emplace_back(box.center_x - vec1_x - vec2_x, box.center_y - vec1_y - vec2_y);
    corners.emplace_back(box.center_x - vec1_x + vec2_x, box.center_y - vec1_y + vec2_y);
​
    return corners;
}
​
​
void* yolo(const char* enginePath,const char* labelPath,float score_thresh,float iou_thresh,int deviceId){
    std::cout<<enginePath<<"\t"<<labelPath<<std::endl;
    auto model = YoloOBB::create_infer(
            std::string(enginePath),
            deviceId,
            score_thresh,
            iou_thresh,
            YoloOBB::NMSMethod::FastGPU,
            500,false
    );
    std::vector<std::string> labels = iLogger::split_string(iLogger::load_text_file(std::string(labelPath)),"\n");
    Yolo* yolo = new Yolo(model,labels);
    return (void*)yolo;
}
​
void release(void* model){
    auto* yolo = (Yolo*)model;
    delete yolo;
    yolo = nullptr;
}
​
char* detect(void* model,void* image,int* count) {
    if(model == nullptr) {
        std::cout<<"model ptr is nullptr \n";
        return nullptr;
    }
    Yolo* yolo = (Yolo*)(model);
    cv::Mat img = *static_cast<cv::Mat*>(image);
    auto objs = yolo->engine->commit(img).get();
    int size = objs.size();
    *count = size;
    if(size > 0) {
        json json_data;
        json_data["objs"] = json::array();
        uint8_t b, g, r;
        for(const auto& obj:objs) {
            json tmp;
            tmp["center_x"] = int(obj.center_x);
            tmp["center_y"] = int(obj.center_y);
            tmp["width"] = int(obj.width);
            tmp["height"] = int(obj.height);
            tmp["angle"] = obj.angle;
            tmp["class_id"] = obj.class_label;
            tmp["conf"] = obj.confidence;
​
            json_data["objs"].push_back(tmp);
            tie(b, g, r) = iLogger::random_color(obj.class_label);
            auto corners = xywhr2xyxyxyxy(obj);
            cv::polylines(img, vector<vector<cv::Point>>{corners}, true, cv::Scalar(b, g, r), 2, 16);
​
            auto name = yolo->labels[obj.class_label];
            auto caption = cv::format("%s %.2f", name.c_str(), obj.confidence);
            int width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10;
            cv::rectangle(img, cv::Point(corners[0].x-3, corners[0].y-33), cv::Point(corners[0].x-3 + width, corners[0].y), cv::Scalar(b, g, r), -1);
            cv::putText(img, caption, cv::Point(corners[0].x-3, corners[0].y-5), 0, 1, cv::Scalar::all(0), 2, 16);
        }
        std::string json_str = json_data.dump();
        int length = json_str.length();
        char* data = (char*)malloc((length + 1) * sizeof(char));
        json_str.copy(data, length,0);
        *(data + length) = '\0';
        return data;
    }
    return nullptr;
}

注意，这里detect返回的是json字符串，同时因为是传入指针，所以会直接修改原始传入图像。gocv本身就是对opencv的一层封装，对于cv::Mat在内存结构上是一致对齐的，因此可以直接转换。而最后返回的字符串指针则需要在go部分使用完后进行释放，否则就会存在内存泄漏。对于之前使用strdup的操作由于是单独开辟一段内存后续没法控制，所以是必然会导致内存泄漏的。

2.2 Go调用编译好的动态库文件

首先放上代码，再来慢慢解释

/*
#cgo LDFLAGS: -L./ -lyolo -lstdc++
#cgo CPPFLAGS: -I /usr/include -I /usr/local/include
#cgo CFLAGS: -std=gnu11
#include<stdio.h>
#include<stdlib.h>
#include "wrap.h"
*/
import "C"
import (
    "fmt"
    "gocv.io/x/gocv"
    "unsafe"
)
​
type Object struct {
    p unsafe.Pointer
}
​
func NewModel(enginePath, labelPath string, score_thresh, iou_thresh float32, deviceId int) *Object {
    path := C.CString(enginePath)
    label_path := C.CString(labelPath)
    obj := &Object{p: C.yolo(path, label_path, C.float(score_thresh), C.float(iou_thresh), C.int(deviceId))}
    C.free(unsafe.Pointer(path))
    C.free(unsafe.Pointer(label_path))
    return obj
}
​
func Release(m *Object) {
    C.release(m.p)
}
​
func Detect(m *Object, img unsafe.Pointer, size *C.int) string {
    res := C.detect(m.p, img, size)
    result := C.GoString(res)
    C.free(unsafe.Pointer(res))
    return result
}
​
func Run() {
    img := gocv.IMRead("../1.jpeg", 1)
    engine_path := "./yolo_obb_dynamic_fp16.engine"
    model := NewModel(engine_path, "./classes.txt", 0.2, 0.5, 0)
    defer Release(model)
​
    var size C.int
    res := Detect(model, unsafe.Pointer(img.Ptr()), &size)
    dst_img := gocv.NewMatFromCMat(unsafe.Pointer(img.Ptr()))
    gocv.IMWrite("./res.jpg", dst_img)
    dst_img.Close()
    fmt.Println(res)
}

最开始注释部分就是cgo的指令，最重要的就是链接到动态库文件。然后就是定义一个Object指针用于存放上面封装的yolo对象。在下面三个函数则是go封装wrap.h中的三个函数，实现对应功能。其中，由于go的string需要转为C.CString才能作为c++函数的传入参数，所以在使用后需要在go部分去释放对应的内存，即C.free(unsafe.Pointer(path))和C.free(unsafe.Pointer(label_path))。如果只是初始化这一次模型，其实也不会额外增加内存占用；如果是多次new推理对象，那就会导致内存增加，不过因为只是文件路径，所以并不会占用太多。C.free(unsafe.Pointer(res))这才是最重要的，因为这是检测后返回的json字符串，占用空间可能会很大（目标很多的情况），并且这段内存已经脱离c++的掌握，这就需要go来free对应的指针。

最后，Run()函数就是利用gocv读取一张图片，创建推理对象并推理。得到结果的同时也得到了绘制结果后的图像指针，利用gocv.NewMatFromCMat(unsafe.Pointer(img.Ptr()))就能将c++的raw ptr重新转为gocv的Mat从而可以保存。那下面就先来看看go调用的效果，为了测试运行时间我稍微改写了一丢丢代码，这里我的测试时间次数为5000次，运行结果如下

image-20241022000629045

image-20241022000324666

image-20241022000717872

3. CGO的服务端方式

想了一下，可能go还是服务端方式居多，大多还是传递base64字符串的方式，所以这一节还是把之前的坑填补一下。

3.1 C++部分的改动

如果是传入base64字符串，我更倾向于在c++中去将base64转为cv::Mat，当然如果你想在go中进行转换，传入gocv的Mat指针也可以，这点就是看个人喜好。

对于base64转为cv::Mat，借助一些base64的库应该不难实现。我这里只给添加的函数及定义

char* infer(void* model,const char* base64,int flag);

char* infer(void* model,const char* base64,int flag) {
    if(model == nullptr) {
        std::cout<<"model ptr is nullptr \n";
        return nullptr;
    }
    Yolo* yolo = (Yolo*)(model);
    std::string base64Img(base64);
    cv::Mat img = Base2Mat(base64Img);
    auto objs = yolo->engine->commit(img).get();
    int size = objs.size();
    if(size > 0) {
        json json_data;
        json_data["objs"] = json::array();
        uint8_t b, g, r;
        for(const auto& obj:objs) {
            json tmp;
            tmp["center_x"] = int(obj.center_x);
            tmp["center_y"] = int(obj.center_y);
            tmp["width"] = int(obj.width);
            tmp["height"] = int(obj.height);
            tmp["angle"] = obj.angle;
            tmp["class_id"] = obj.class_label;
            tmp["conf"] = obj.confidence;
​
            json_data["objs"].push_back(tmp);
            if(flag == 1){
                tie(b, g, r) = iLogger::random_color(obj.class_label);
                auto corners = xywhr2xyxyxyxy(obj);
                cv::polylines(img, vector<vector<cv::Point>>{corners}, true, cv::Scalar(b, g, r), 2, 16);
​
                auto name = yolo->labels[obj.class_label];
                auto caption = cv::format("%s %.2f", name.c_str(), obj.confidence);
                int width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10;
                cv::rectangle(img, cv::Point(corners[0].x-3, corners[0].y-33), cv::Point(corners[0].x-3 + width, corners[0].y), cv::Scalar(b, g, r), -1);
                cv::putText(img, caption, cv::Point(corners[0].x-3, corners[0].y-5), 0, 1, cv::Scalar::all(0), 2, 16);
            }
        }
        if(flag == 1){
            json_data["dst_img"] = Mat2Base64(img,"jpg");
        }
        std::string json_str = json_data.dump();
        int length = json_str.length();
        char* data = (char*)malloc((length + 1) * sizeof(char));
        json_str.copy(data, length,0);
        *(data + length) = '\0';
        return data;
    }
    return nullptr;
}

没有过多的改变，只是添加了base64与Mat相互转换的代码。

3.2 Go搭建一个简单服务

func Infer(m *Object, img string, flag int) string {
    base64 := C.CString(img)
    res := C.infer(m.p, base64, C.int(flag))
    result := C.GoString(res)
    C.free(unsafe.Pointer(base64))
    C.free(unsafe.Pointer(res))
    return result
}
​

这就是对应的go部分的封装，然后利用echo搭建一个简单的服务

type Param struct {
    Image string `json:"image"`
    Flag  int    `json:"flag"`
}
​
func main() {
    engine_path := "./yolo_obb_dynamic_fp16.engine"
    model := infer.NewModel(engine_path, "./cmd/classes.txt", 0.2, 0.5, 0)
    defer infer.Release(model)
    e := echo.New()
    var start time.Time
    //param := new(Param)
    e.POST("/obb", func(c echo.Context) error {
        param := new(Param)
        if err := c.Bind(param); err != nil {
            return c.String(http.StatusBadRequest, "{'msg','invalid param!'}")
        }
        start = time.Now()
        result := infer.Infer(model, param.Image, param.Flag)
        log.Printf("time: %s,draw flag: %v infer cost: %s \n", start, param.Flag, time.Since(start))
        return c.String(http.StatusOK, result)
    })
​
    log.Fatal(e.Start(":1323"))
}

传入的图片如下

image-20241022004307943

首先不绘制结果图，即flag不为1来看看响应速度

image-20241022004449168

image-20241022004502019

推理部分耗时大概在8.5ms左右，整体耗时在15ms以内。添加绘制后耗时会有所增加，推理耗时提升到15ms左右，主要还是检测到的目标多遍历绘制会导致耗时增加，绘制结果如下

image-20241022004634504

image-20241022004831642

image-20241022004813258

3.3 压测来看看内存情况

利用apifox等工具可以比较简单的进行压测，但是如果是本地压测为了看清我们的服务内存变化，可以使用top -p pid来观察。对于内存泄漏，当然是在越小内存的机器上表现得越明显，所以不妨写好后在边缘端设备试试。

image-20241022005815316

压测期间稳定在16.398G左右，RES也没有异常升高

image-20241022010027724

avatar
文章
阅读
粉丝