1. Johngo学长首页
  2. 技术点详细复盘专题
  3. 人工智能

【TensorRT】记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程

• 人工智能 • 阅读 51

记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程

最近因为课题的原因，需要部署下YOLOv5的模型。之前一般部署YOLOv5的常规方法是直接使用Wangxinyu大佬的tensorrtx这个仓库去部署，因为之前的YOLOv5转trt真的非常费劲。现在YOLOv5推出了v6.1之后，支持直接使用官方repo里面的 export.py脚本直接导出trt的engine，对部署党来说真的是喜大普奔。因此在捣鼓了半天之后，成功使用C++完成对YOLOv5模型的部署，正好最近看有同学在问这个事儿，因此在此记录下来。

1. 导出trt engine

首先我们按照YOLOv5官方的Installation配置好环境，下载yolov5s.pt权重之后，在命令行运行 export.py脚本：

python3 export.py --weights ./yolov5s.pt --include engine --imgsz 640 --device 0

解释下其中几个命令行参数：

  • –weights：YOLOv5的权重路径；
  • –include：需要将PyTorch模型转换成什么格式；
  • –imgsz：输入模型大小；
  • –device：在什么设备商运行。因为TensorRT跟硬件强相关，因此需要指定你使用的是机器里面的哪块卡。

经过一番操作之后，可以看到成功在路径下导出了 yolov5s.engine文件，这就是TensorRT的推理引擎。

2. 测试TensorRT推理引擎是否可用

YOLOv5的官方repo提供了 detect.py脚本，可以供我们测试模型权重是否可用。我们在命令行运行 detect.py脚本：

python detect.py --weights yolov5s.engine --imgsz 640 --device 0

运行后命令行会输出如下的信息，可以看到，推理只需要0.003s即可完成，即3ms，与使用PyTorch推理需要的13ms相比提升了很多。

【TensorRT】记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程

本来至此应该推理也就结束了，但是工业场景往往Python并不合适，因此我们使用C++完成推理部分。

3. C++部署TensorRT

(1) 初始化TensorRT引擎

在进行推理之前，首先需要先初始化TensorRT的引擎。代码如下：

void YOLOv5::initialize()
{
    cudaSetDevice(0);
    char *trtModelStream{nullptr};
    size_t size{0};

    std::ifstream file(mEnginePath, std::ios::binary);
    std::cout << "[I] Detection model creating...\n";
    if (file.good())
    {
        file.seekg(0, file.end);
        size = file.tellg();
        file.seekg(0, file.beg);
        trtModelStream = new char[size];
        assert(trtModelStream);
        file.read(trtModelStream, size);
        file.close();
    }

    mRuntime = createInferRuntime(mGLogger);
    assert(mRuntime != nullptr);

    std::cout << "[I] Detection engine creating...\n";
    mEngine = mRuntime->deserializeCudaEngine(trtModelStream, size);
    assert(mEngine != nullptr);
    mContext = mEngine->createExecutionContext();
    assert(mContext != nullptr);
    delete[] trtModelStream;

    auto out_dims = mEngine->getBindingDimensions(1);

    mBlob = new float[mInputSize];
    mProb = new float[mOutputSize];

    assert(mEngine->getNbBindings() == 2);
    std::cout << "[I] Cuda buffer creating...\n";

    mInputIndex = mEngine->getBindingIndex("images");

    assert(mEngine->getBindingDataType(mInputIndex) == nvinfer1::DataType::kFLOAT);
    mOutputIndex = mEngine->getBindingIndex("outputs");
    assert(mEngine->getBindingDataType(mOutputIndex) == nvinfer1::DataType::kFLOAT);

    checkStatus(cudaMalloc(&mBuffers[mInputIndex], mInputSize * sizeof(float)));
    checkStatus(cudaMalloc(&mBuffers[mOutputIndex], mOutputSize * sizeof(float)));

    std::cout << "[I] Cuda stream creating...\n";
    checkStatus(cudaStreamCreate(&mStream));

    std::cout << "[I] Detection engine created!\n";
}

需要注意的是其中几个跟待部署模型强相关的参数：

  • mInputSize：模型输入的大小，由于我们这个场景一次输入进一张图像即可，yolov5s的输入就是 3 * 640 * 640，因此我们的输入尺寸也应该是 1 * 3 * 640 * 640，即我们要创建一个这么大的一维数组；
  • mOutputSize：模型输出的大小，YOLOv5的输出头一共有三个，每个输出头的大小分别为 20 * 20、 40 * 40、 80 * 80，一共是有8400个栅格(grid)，每个栅格有3个anchor，每个anchor会输出85个信息（80个类别 + xywh + confidence）。YOLOv5的作者为了方便我们做后处理，因此在模型导出的时候，将三个输出头Concat在了一起，形成了 1 * 25200 * 85大小的输出。其中25200就是 20 * 20 + 40 * 40 + 80 * 80。

【TensorRT】记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程
  • mInputIndex和mOutputIndex：这两个要通过 mEngine->getBindingIndex来获得索引，这个方法输入节点的name后可以返回索引，因此我们要清楚输入和输出的name分别是什么，用Netron打开ONNX看看就知道了，或者看看 export.py的 export_onnx函数看看设置的 input_names和 output_names分别是什么；

【TensorRT】记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程

(2) 预处理图像

图像在输入之前肯定要先做预处理，主要就是做resize和normalize。

void YOLOv5::preprocess(cv::Mat& src, cv::Mat& dst)
{
    mCvOriginSize = src.size();
    dst = src.clone();
    cv::cvtColor(dst, dst, cv::COLOR_BGR2RGB);
    cv::resize(dst, dst, mCvInputSize);
    dst.convertTo(dst, CV_32FC3);

    dst = dst / 255.0f;
}

非常简单，一共就这么几步：

  • 通道顺序从BGR转为RGB（OpenCV默认输入图像后通道顺序是BGR）；
  • resize到640；
  • 把矩阵转为Float32型（不然除以255可能会出问题）；
  • normalize（除以255）。

(3) 将输入的图像的每一个像素按顺序存入数组

TensorRT并不能够直接以OpenCV的Mat数据结构为输入，需要我们先将Mat里面的每一个像素存进数组内。我们先前已经声明了两个成员变量 mBlob和 mProb，我们现在就要将输入数据存入 mBlob中。

void YOLOv5::blobFromImage(cv::Mat& img)
{
    preprocess(img, img);
    int channels = img.channels();
    int cols = img.cols;
    int rows = img.rows;

    for (int c = 0; c < channels; c++)
    {
        for (int row = 0; row < rows; row++)
        {
            for (int col = 0; col < cols; col++)
            {
                mBlob[c * rows * cols + row * cols + col] = img.at<cv::Vec3f>(row, col)[c];
            }
        }
    }
}

预处理之后，按照一行一行的顺序把图像的像素存入 mBlob中即可，这一步也很简单。

(4) 执行推理步骤

int YOLOv5::doInference()
{
    checkStatus(cudaMemcpyAsync(mBuffers[mInputIndex], mBlob, mInputSize * sizeof(float), cudaMemcpyHostToDevice, mStream));
    mContext->enqueueV2(mBuffers, mStream, nullptr);
    checkStatus(cudaMemcpyAsync(mProb, mBuffers[mOutputIndex], mOutputSize * sizeof(float), cudaMemcpyDeviceToHost, mStream));
    cudaStreamSynchronize(mStream);

    return 0;
}

TensorRT的推理非常精简，一共就是三步走：

  1. 将 mBlob里面的数据拷贝至显卡；
  2. 用 mContext->enqueueV2方法执行推理；
  3. 将推理的结果拷贝至内存（也就是 mProb里）。

(5) 后处理

这一步才是精髓。我们得到的结果应该包括以下部分：

  1. 类别概率，有25200 * 80个；
  2. 位置信息，是相对于Anchor和Grid的偏移量，有25200 * 4个；
  3. 置信度，有25200 * 1个；

但是YOLOv5团队为了方便我们做后处理，已经将第2点的位置信息的解码过程一并导出到了ONNX中，自然也随着ONNX一并转到了TensorRT里面。也就是说，Engine推理后输出的位置信息，就是真实的位置信息（相对于640 * 640而言），不需要我们再费劲写位置信息的解码过程。

我们再明确下一共要干哪几件事情：

  1. 整理输出结果，置信度低于置信度阈值的不保留；
  2. 做NMS；

我们先做第一件事情：

struct Object
{
    cv::Rect rect;
    int label;
    float conf;
};

std::pair<int, float> YOLOv5::argmax(std::vector<float>& vSingleProbs)
{
    std::pair<int, float> result;
    auto iter = std::max_element(vSingleProbs.begin(), vSingleProbs.end());
    result.first = static_cast<int>(iter - vSingleProbs.begin());
    result.second = *iter;

    return result;
}

void YOLOv5::generate_proposals(std::vector<Object>& objects, float confThresh)
{
    int nc = 80;
    for (int i = 0; i < 25200; i++)
    {
        float conf = mProb[i * (nc + 5) + 4];
        if (conf > confThresh)
        {
            Object obj;
            float cx = mProb[i * (nc + 5)];
            float cy = mProb[i * (nc + 5) + 1];
            float w  = mProb[i * (nc + 5) + 2];
            float h  = mProb[i * (nc + 5) + 3];
            obj.rect.x = static_cast<int>(cx - w * 0.5f);
            obj.rect.y = static_cast<int>(cy - h * 0.5f);
            obj.rect.width = static_cast<int>(w);
            obj.rect.height = static_cast<int>(h);

            std::vector<float> vSingleProbs(nc);
            for (int j = 0; j < vSingleProbs.size(); j++)
            {
                vSingleProbs[j] = mProb[i * 85 + 5 + j];
            }

            auto max = argmax(vSingleProbs);
            obj.label = max.first;
            obj.conf = conf;

            objects.push_back(obj);
        }
    }
}

可以看到在generate_proposals函数内，我们做了这样的几件事情。

  1. 遍历所有结果，先取排在下标为4的置信度（顺序是x y w h conf），判断是否高于置信度的阈值；
  2. 如果高于阈值，按照顺序取xywh（注意是xy是中心点坐标，但是cv::Rect的xy是左上角点坐标）；
  3. 将xywh整理进cv::Rect数据结构内；
  4. 用argmax方法从后80个数据内获得类别的label，不需要多解释；

这样就完成了第一步的整理。但此时我们的框会有很多冗余，这在目前常用的目标检测算法里面非常常见。因为每一个Grid和每一个Anchor都会输出一个结果，而目标附近的Grid和Anchor输出的结果很大可能指的都是同一个目标，因此就会出现目标处会有很多框重叠在一起的情况。这就需要用nms算法去把框筛一下。

void YOLOv5::qsort_descent_inplace(std::vector<Object>& objects, int left, int right)
{
    int i = left;
    int j = right;
    float p = objects[(left + right) / 2].conf;

    while (i  j)
    {
        while (objects[i].conf > p)
            i++;

        while (objects[j].conf < p)
            j--;

        if (i  j)
        {

            std::swap(objects[i], objects[j]);

            i++;
            j--;
        }
    }

    #pragma omp parallel sections
    {
        #pragma omp section
        {
            if (left < j) qsort_descent_inplace(objects, left, j);
        }
        #pragma omp section
        {
            if (i < right) qsort_descent_inplace(objects, i, right);
        }
    }
}

void YOLOv5::qsort_descent_inplace(std::vector<Object>& objects)
{
    if (objects.empty())
        return;

    qsort_descent_inplace(objects, 0, objects.size() - 1);
}

void YOLOv5::nms_sorted_bboxes(const std::vector<Object>& vObjects, std::vector<int>& picked, float nms_threshold)
{
    picked.clear();

    const int n = vObjects.size();

    std::vector<float> areas(n);
    for (int i = 0; i < n; i++)
    {
        areas[i] = vObjects[i].rect.area();
    }

    for (int i = 0; i < n; i++)
    {
        const Object& a = vObjects[i];

        int keep = 1;
        for (int j = 0; j < (int)picked.size(); j++)
        {
            const Object& b = vObjects[picked[j]];

            float inter_area = intersection_area(a, b);
            float union_area = areas[i] + areas[picked[j]] - inter_area;

            if (inter_area / union_area > nms_threshold)
                keep = 0;
        }

        if (keep)
            picked.push_back(i);
    }
}

此处做nms的方法是嫖的NCNN的solution，ncnn yyds！

我们将逻辑整理一下，形成 decodeOutputs方法：

std::vector<Object> YOLOv5::decodeOutputs(std::vector<Object>& objects)
{
    generate_proposals(objects, 0.2f);
    qsort_descent_inplace(objects);

    std::vector<int> picked;
    nms_sorted_bboxes(objects, picked, 0.45f);

    int count = picked.size();

    int img_w = mCvOriginSize.width;
    int img_h = mCvOriginSize.height;
    float scaleH =  static_cast<float>(mCvInputSize.height) / static_cast<float>(img_h);
    float scaleW =  static_cast<float>(mCvInputSize.width) / static_cast<float>(img_w);

    std::vector<Object> results;
    results.resize(count);
    for (int i = 0; i < count; i++)
    {
        Object obj = objects[picked[i]];

        float x0 = static_cast<float>(obj.rect.x) / scaleW;
        float y0 = static_cast<float>(obj.rect.y) / scaleH;
        float x1 = static_cast<float>(obj.rect.x + obj.rect.width) / scaleW;
        float y1 = static_cast<float>(obj.rect.y + obj.rect.height) / scaleH;

        x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.0f);
        y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.0f);
        x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.0f);
        y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.0f);

        obj.rect.x = static_cast<int>(x0);
        obj.rect.y = static_cast<int>(y0);
        obj.rect.width = static_cast<int>(x1 - x0);
        obj.rect.height = static_cast<int>(y1 - y0);
        results[i] = obj;
    }

    return results;
}

从代码中可以看到，首先调用了 generate_proposals进行低置信度目标的过滤，接着调用 qsort_descent_inplace方法做快速排序，再调用 nms_sorted_bboxes方法做nms。然后我们获得640与原图的宽高的比例，将框映射回原图，最后纠正下过大的框和负数坐标即可。

至此YOLOv5的TensorRT部署也就结束了，我们可以画个框看看：

【TensorRT】记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程

Original: https://blog.csdn.net/ChunjieShan/article/details/126120013
Author: 单胖
Title: 【TensorRT】记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程

原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/707084/

转载文章受原作者版权保护。转载请注明原作者出处！

赞 (0)
0

【自取】最近整理的，有需要可以领取学习：



Linux核心资料大放送~

全栈面试题汇总（持续更新&可下载）

一个提高学习100%效率的工具！

【超详细】深度学习面试题目！

LeetCode Python刷题答案下载！

LeetCode Java版刷题答案下载！

LeetCode C++ 版本，抓紧保存！

LeetCode GO语言 刷题答案下载！

大家都在看

亲爱的 Coder【最近整理，可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球