重写TensorRT版本YOLOX部署代码

一天到晚潜水的鱼

前言

YOLOX 是前段时间旷视开源的一个目标检测算法，据说效果很好，这两天有空了就准备研究一下，看了论文感觉里面干货还是很多的，等后面再仔细研究研究。从论文放出的结果来看， YOLOX 在速度和精度上应该是全面超过了之前的 YOLO 系列算法的。

比较良心的是，作者不仅开源了代码和模型，还放出了 TensorRT、OpenVINO、NCNN 等框架下的模型部署示例代码，可谓是工程人的福音。

看了 TensorRT 版本的 C++ 部署示例代码，决定自己重新写一下，就当练手了。

实现过程

这里主要记录需要注意的事项和与官方示例代码不一样的地方。

1. 下载ONNX模型

ONNX 模型可以从下面的链接页面中下载:

https://github.com/Megvii-BaseDetection/YOLOX/releases/

需要注意的是，我们需要下载 0.1.1pre 版本的权重，最新的代码中作者已经修改了图像预处理的方式，这会导致之前版本的 ONNX 模型权重与最新的代码不兼容，这是作者的说明：

2. TensorRT解析ONNX模型

YOLOX 官方提供的 TensorRT 版本示例代码是先通过 tools/trt.py 脚本把 ONNX 模型解析再后序列化到 model_trt.engine 文件中，然后 C++ 代码再从该文件中加载模型去进行推理。这里我们可以直接在 C++ 代码中去解析 ONNX 模型，然后把它序列化到 .engine 文件中， TensorRT 解析 ONNX 模型的方法可以参考英伟达官方提供的 sampleOnnxMNIST 例程。

if (!isFileExists(engine_path)) {
  std::cout << "The engine file " << engine_path
            << " has not been generated, try to generate..." << std::endl;
  engine_ = SerializeToEngineFile(model_path_, engine_path);
  std::cout << "Succeed to generate engine file: " << engine_path
            << std::endl;
} else {
  std::cout << "Use the exists engine file: " << engine_path << std::endl;
  engine_ = LoadFromEngineFile(engine_path);

这里首先判断 ONNX 模型对应的 .engine 文件是否存在，如果存在就直接从 .engine 文件中加载模型，否则就创建一个 ONNX 模型解析器去解析模型，然后把模型序列化到 .engine 文件中方便下次使用。

//把模型序列化到engine文件中
nvinfer1::IHostMemory *trtModelStream = engine->serialize();
std::stringstream gieModelStream;
gieModelStream.seekg(0, gieModelStream.beg);
gieModelStream.write(static_cast<const char *>(trtModelStream->data()),
                      trtModelStream->size());
std::ofstream outFile;
outFile.open(engine_path);
outFile << gieModelStream.rdbuf();
outFile.close();

3. 自动获取模型输入尺寸

官方示例代码中，默认模型的输入尺寸是 640x640

static const int INPUT_W = 640;
static const int INPUT_H = 640;

但是如果模型的输入尺寸是 416x416 或者是长宽不等的 512x416 这种尺寸，那么就还需要改代码，感觉不是很方便。其实模型的输入输出维度都可以通过 TensorRT 提供的接口获取，这样就方便根据模型解析的结果自动获取输入尺寸，然后根据这个信息去对输入图像做 resize 了。

nvinfer1::Dims input_dim = engine_->getBindingDimensions(index);
int input_size = 1;
for (int j = 0; j < input_dim.nbDims; ++j) {
  input_size *= input_dim.d[j];

上面的代码中， input_dim.d 为模型的输入尺寸，按照 NCHW 的顺序排列。

4. 图像预处理

官方示例代码中，预处理的时候是对图像做长宽等比例缩放，不足的地方再进行填充：

cv::Mat static_resize(cv::Mat& img) {
  float r = std::min(INPUT_W / (img.cols*1.0), INPUT_H / (img.rows*1.0));
  int unpad_w = r * img.cols;
  int unpad_h = r * img.rows;
  cv::Mat re(unpad_h, unpad_w, CV_8UC3);
  cv::resize(img, re, re.size());
  cv::Mat out(INPUT_H, INPUT_W, CV_8UC3, cv::Scalar(114, 114, 114));
  re.copyTo(out(cv::Rect(0, 0, re.cols, re.rows)));
  return out;

我就直接简单粗暴地 resize 了(不要学我)：

cv::Mat resize_image;
cv::resize(input_image, resize_image, cv::Size(model_width_, model_height_));

两种方法的对比：

5. 后处理

后处理是对模型推理的结果进行解析， YOLOX 是 anchor-free 的目标检测算法，解析的时候相对要简单一些。与 YOLOv3 类似， YOLOX 还是在3个尺度上去做检测，每一层特征图上的单元格只预测一个框,每个单元格输出的内容是 x,y,w,h,objectness 这5个内容再加上每个类别的概率。可以用 Netron 看一下模型后面几层的结构:

可以看到，如果模型输入尺寸为 640x640 ，分别降采样 8,16,32 倍后得到的特征图尺寸分别为 80x80,40x40,20x20 ， COCO 数据集有 80 个类别那么每个特征图的单元格输出的数据长度为 5+80=85 ，3个特征图上的结果最终会 concat 到一起进行输出，所以最终输出的数据维度为 (80x80+40x40+20x20)x85=8400x85 。

官方示例代码中用了好几个函数做后处理，感觉有点繁琐，于是我重写了这部分的代码：

const std::vector<int> strides = {8, 16, 32};
float *ptr = const_cast<float *>(output);
for (std::size_t i = 0; i < strides.size(); ++i) {
  const int stride = strides.at(i);
  const int grid_width = model_width_ / stride;
  const int grid_height = model_height_ / stride;
  const int grid_size = grid_width * grid_height;
  for (int j = 0; j < grid_size; ++j) {
    const int row = j / grid_width;
    const int col = j % grid_width;
    const int base_pos = j * (kNumClasses + 5);
    const int class_pos = base_pos + 5;
    const float objectness = ptr[base_pos + 4];
    const int label =
        std::max_element(ptr + class_pos, ptr + class_pos + kNumClasses) -
        (ptr + class_pos);
    const float confidence = (*(ptr + class_pos + label)) * objectness;
    if (confidence > confidence_thresh) {
      const float x = (ptr[base_pos + 0] + col) * stride / width_scale;
      const float y = (ptr[base_pos + 1] + row) * stride / height_scale;
      const float w = std::exp(ptr[base_pos + 2]) * stride / width_scale;
      const float h = std::exp(ptr[base_pos + 3]) * stride / height_scale;
      Object obj;
      obj.box.x = x - w * 0.5f;
      obj.box.y = y - h * 0.5f;
      obj.box.width = w;
      obj.box.height = h;
      obj.label = label;
      obj.confidence = confidence;