[Python知识库]【TensorRT】记一次使用C++接口TensorRT部署yolov5 v6.1模型的过程

记一次使用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相比提升了很多。

本来至此应该推理也就结束了，但是工业场景往往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];

    // Pointers to input and output device buffers to pass to engine.
    // Engine requires exactly IEngine::getNbBindings() number of buffers.
    assert(mEngine->getNbBindings() == 2);
    std::cout << "[I] Cuda buffer creating...\n";

    // In order to bind the buffers, we need to know the names of the input and output tensors.
    // Note that indices are guaranteed to be less than IEngine::getNbBindings()
    mInputIndex = mEngine->getBindingIndex("images");

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

    // Create GPU buffers on device
    checkStatus(cudaMalloc(&mBuffers[mInputIndex], mInputSize * sizeof(float)));
    checkStatus(cudaMalloc(&mBuffers[mOutputIndex], mOutputSize * sizeof(float)));

    // Create stream
    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。

• mInputIndex和mOutputIndex：这两个要通过mEngine->getBindingIndex来获得索引，这个方法输入节点的name后可以返回索引，因此我们要清楚输入和输出的name分别是什么，用Netron打开ONNX看看就知道了，或者看看export.py的export_onnx函数看看设置的input_names和output_names分别是什么；

(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)
        {
            // swap
            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]];

            // intersection over union
            float inter_area = intersection_area(a, b);
            float union_area = areas[i] + areas[picked[j]] - inter_area;
            // float IoU = inter_area / union_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]];

        // adjust offset to original unpadded
        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;

        // clip
        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部署也就结束了，我们可以画个框看看：