【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的推理非常精简，一共就是三步走：
- 将
mBlob里面的数据拷贝至显卡； - 用
mContext->enqueueV2方法执行推理； - 将推理的结果拷贝至内存（也就是
mProb里）。
(5) 后处理
这一步才是精髓。我们得到的结果应该包括以下部分：
- 类别概率，有25200 * 80个；
- 位置信息，是相对于Anchor和Grid的偏移量，有25200 * 4个；
- 置信度，有25200 * 1个；
但是YOLOv5团队为了方便我们做后处理，已经将第2点的位置信息的解码过程一并导出到了ONNX中，自然也随着ONNX一并转到了TensorRT里面。也就是说，Engine推理后输出的位置信息，就是真实的位置信息（相对于640 * 640而言），不需要我们再费劲写位置信息的解码过程。
我们再明确下一共要干哪几件事情：
- 整理输出结果，置信度低于置信度阈值的不保留；
- 做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函数内，我们做了这样的几件事情。
- 遍历所有结果，先取排在下标为4的置信度（顺序是x y w h conf），判断是否高于置信度的阈值；
- 如果高于阈值，按照顺序取xywh（注意是xy是中心点坐标，但是cv::Rect的xy是左上角点坐标）；
- 将xywh整理进cv::Rect数据结构内；
- 用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部署也就结束了，我们可以画个框看看：