如何用TensorRT部署YOLOv6

本文首发于微信公众号【DeepDriving】,关注公众号后台回复关键字【YOLOv6】可获取本文代码链接。

前言

YOLOv6是美团视觉智能部研发的一个致力于工业应用的目标检测算法,该算法框架同时专注于检测的精度和推理效率。在官方发布的文章中,宣称YOLOv6的精度与速度都远超YOLOv5YOLOX。在部署方面,YOLOv6支持GPU(TensorRT)、CPU(OPENVINO)、ARM(MNN、TNN、NCNN)等不同平台的部署,极大地简化工程部署时的适配工作。

在这里插入图片描述

YOLOv6具体的实现细节大家可以去看官方的文章,这里我就不做过多的介绍了。本文主要介绍如何用TensorRT部署YOLOv6C++实现)。

实现过程

1. 下载ONNX模型

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

https://github.com/meituan/YOLOv6/releases/tag/0.1.0

也可以下载PyTorch格式的模型文件,然后用官方提供的脚本转换为ONNX模型:

python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1

上面两种方式得到的ONNX模型最好用onnx-simplifier工具再处理一下,这样得到的模型会更加精简,看上去会舒服很多。

import onnx
from onnxsim import simplify

model = onnx.load('yolov6s.onnx')
model_simple, check = simplify(model)
assert check, 'Failed to simplify model'
onnx.save(model_simple, 'yolov6s_simplify.onnx')
print('Succeed to simplify model')

2. TensorRT解析ONNX模型

这部分的代码跟我之前部署YOLOX的代码是一样的,感兴趣的可以参考我之前写的文章

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

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);
}

3. 图像预处理

这里与官方的处理不一样,我做预处理的时候没有对图像去做等比例缩放然后不足的地方再进行填充,而是直接做resize了:

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

两种预处理方法的对比:

在这里插入图片描述

在这里插入图片描述

可以看到,直接Resize会导致图像中的物体变形,所以还是不建议这么做,我这么做是因为比较懒。

Resize操作后,需要对图像的每个像素除以255进行归一化,图像数据在内存中要按照CHW的顺序进行排列。

5. 后处理

YOLOX一样,YOLOv6是一个anchor-free的目标检测算法,模型还是在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

需要注意的是,上图红框中的算子是在做后处理操作,就是把x,y,w,h做相应的操作然后还原回相对于模型输入尺寸的大小。既然这些操作都在推理的时候做了,那我们在得到模型的推理结果后,就只需要进行简单的处理就可以了:

float *ptr = const_cast<float *>(output);
for (int i = 0; i < 8400; ++i, ptr += (kNumClasses + 5)) {
  const float objectness = ptr[4];
  if (objectness >= kObjectnessThresh) {
    const int label =
        std::max_element(ptr + 5, ptr + (kNumClasses + 5)) - (ptr + 5);
    const float confidence = ptr[5 + label] * objectness;
    if (confidence >= confidence_thresh) {
      const float bx = (ptr[0]);
      const float by = (ptr[1]);
      const float bw = ptr[2];
      const float bh = ptr[3];

      Object obj;
      obj.box.x = (bx - bw * 0.5f) / width_scale;
      obj.box.y = (by - bh * 0.5f) / height_scale;
      obj.box.width = bw / width_scale;
      obj.box.height = bh / height_scale;
      obj.label = label;
      obj.confidence = confidence;
      objs->push_back(std::move(obj));
    }
  }
}  

最后,需要对模型输出的结果做非极大值抑制操作以去除重复的框,我是用Soft-NMS做的。

结果

yolov6_s.onnx模型测试的几个结果如下:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

在我的GeForce GTX 1050 Ti显卡上测试各个模型的耗时如下表所示:

模型输入尺寸耗时
yolov6n.onnx640x6408 ms
yolov6t.onnx640x64021 ms
yolov6s.onnx640x64023 ms

顺便把YOLOX的耗时也贴一下:

模型输入尺寸耗时
yolox_nano.onnx640x6408 ms
yolox_tiny.onnx640x64013 ms
yolox_s.onnx640x64018 ms
yolox_m.onnx640x64039 ms
yolox_l.onnx640x64075 ms

总结

YOLOv6YOLOX的基础上针对BackBone、Neck、Head以及训练策略等方面做了一些改进,效果还是非常不错的,但是个人觉得还没达到“远超”的水平。对于我这种白嫖党来说,非常乐意看到YOLOX或者YOLOv6这样效果好还容易部署的算法开源出来,真心希望这样的工作更多一些,哈哈…

  • 1
    点赞
  • 8
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 2
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

DeepDriving

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值