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. 后处理

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

可以看到,如果模型输入尺寸为640x640,分别降采样8,16,32倍后得到的特征图尺寸分别为80x80,40x40,20x20COCO数据集有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;
      objs->push_back(std::move(obj));
  ptr += grid_size * (kNumClasses + 5);

这里有个小技巧:没有必要每个单元格都先去解析出x,y,w,h后再去看置信度是否大于阈值,而应该先判断置信度,如果置信度大于阈值再去解析x,y,w,h。这样做还是可以减少很多计算量的(在嵌入式平台上还是能省则省吧),毕竟一般几千个单元格的结果可能只有几十个是符合条件的。

最后我用Soft-NMS算法(不知道我写的对不对)做非极大值抑制去除重复的框。

yolox_s.onnx模型测试的几个结果:

GeForce RTX 2080显卡上测试的的各个模型的耗时如下表所示(不是很精确):

模型输入尺寸耗时
yolox_darknet.onnx640x64022 ms
yolox_l.onnx640x64021 ms
yolox_m.onnx640x64011 ms
yolox_s.onnx640x6407 ms
yolox_tiny.onnx416x4163 ms
yolox_nano.onnx416x4162 ms

针对上一节中的第3点,我试了不管模型输入尺寸是640x640416x416还是512x416都是可以的,程序会做自适应处理。

一句话:YOLOX又快又好!

欢迎关注我的公众号【DeepDriving】,我会不定期分享计算机视觉、机器学习、深度学习、无人驾驶等领域的文章。

抵制代码重写昨天,一位老上级邀请我一起吃午餐。当坐在哪里等待上菜时,我们缅怀起早期这个公司的往事。他有一句话让我心里一虚:啊,你这个判官…我记得当你看到Dan(公司的第一位程序员)写的代码时的样子。你说:“这代码写的真烂,需要重写!”   昨天,一位老上级邀请我一起吃午餐。当坐在哪里等待上菜时,我们缅怀起早期这个公司的往事。他有一句话让我心里一虚:   啊,你这个判官…我记得当你看到Dan(公司的第一位程序员)写的代码时的样子。你说:“这代码写的真烂,需要重写!”   我恐怕是没有足够的勇气告诉他,我这“代码需要重写”的主张是错误的。不错,我认为这代码写的很乱。但是,据过去历次的经验,我感觉大
基本思想:之前很多项目都是基于TensorRT部署nano,今天顺手记录一下tengine部署nano吧~ 分了四步走:先进行mingw32的交叉编译,以运行在clion上,官方似乎没有提供window10+mingw32的教程,并进行验证 :https://tengine-docs.readthedocs.io/zh_CN/latest/index.html 进行nano的tensorRT部署,...
YOLOX Window10 TensorRT 全面部署教程配置流程训练自己的模型安装torch2trt转换模型编译DemoWindow10 部署Window10上使用自己训练的模型进行tensorrt检测总结 官网这部分写的非常笼统,我把完整教程分享给大家。 window10+GTX1650+TensorRT-7.2.3.4-CUDA11.1 yolox程序在ubuntu上,所以部分内容是ubuntu上操作 其他环境不保证 训练自己的模型 训练部分请参照官网流程进行训练,或者直接使用官方
yolo系列是目标识别的重头戏,为了更好的理解掌握它,我们必须从源码出发深刻理解代码。下面我们来讲解pytorch实现的yolov3源码。在讲解之前,大家应该具备相应的原理知识yolov1,yolov2,yolov3。 大部分同学在看论文时并不能把所有的知识全部掌握。我们必须结合代码代码将理论变成实践),它是百分百还原理论的,也只有在掌握代码以及理论后,我们才能推陈出新有所收获,所以大家平时一定...
Person person = (Person) obj; return person.getName().equals(name) && person.getAge() == age; 上述代码中,我们通过重写equals方法来比较两个Person对象是否相等。我们首先判断两个对象是否为同一个引用,如果是则直接返回true。然后我们判断obj是否是Person类型的对象,如果不是则返回false。最后我们将obj强制转换为Person类型的对象,比较其name和age属性是否与当前对象的相同,如果相同则返回true,否则返回false。 AI小花猫: 在windows上也是一样的吗? 我这边直接这个命令报错了:报错如下:RROR: '.[dev]' is not a valid editable requirement. It should either be a path to a local project or a VCS URL (beginning with bzr+http, bzr+https, bzr+ssh, bzr+sftp, bzr+ftp, bzr+lp, bzr+file, git+http, git+https, git+ssh, git+git, git+file, hg+file, hg+http, hg+https, hg+ssh, hg+static-http, svn+ssh, svn+http, svn+https, svn+svn, svn+file). YOLOv8初体验:检测、跟踪、模型部署 DeepDriving: 直接复制这行命令执行就好了