Yolov5 不做赘述,目前目标检测里使用非常多的模型,效果和速度兼顾,性能强悍,配合TensorRT推理加速,在工业界可以说是非常流行的组合。
废话不多说,直接开整,以下使用的Tensor RT部署推理路线为:Pytorch-> ONNX -> TensorRT。
pytorch导出到onnx模型,可以非常方便,并且支持dynamic维度,配合netron工具,可以查看模型的网络结构,而TensorRT对ONNX的支持也非常完整,所以选择这一套流程,可以非常轻松的完成TensorRT的部署,同时,tensor RT提供官方的nms插件,使得推理代码可以免去编写nms的部分,极大提高效率。
GitHub仓库地址: gentlemanarch/yolov5-tensorrt
环境准备
-
系统:Ubuntu20.04LTS系统,或者tensorRT官方docker镜像:nvcr.io/nvidia/tensorrt:21.05-py3(推荐)
-
TensorRT: 本篇使用的TensorRT7.2.3
-
Yolov5: 截止2021.09.11的develop分支代码
-
gcc: 9.3.0
-
torch: 1.8.2
-
onnx:1.10.1
-
onnx-simplifier:0.3.6
pytorch模型训练
clone Yolov5的官方代码,按照教程训练得到pt权重文件。
github地址:
https://github.com/ultralytics/yolov5
Torch -> onnx
在导出到onnx之前,为了方便后续添加nms插件,需要对torch的模型输出做一些修改.
models/yolo.py:
将这部分代码
if self.inplace:
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no))
替换为:
if self.inplace:
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2).expand(bs, self.na, 1, 1, 2) # wh
rest = y[..., 4:]
yy = torch.cat((xy, wh, rest), -1)
z.append(yy.view(bs, -1, self.no))
else:
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].view(1, self.na, 1, 1, 2) # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
同时,这句:
return x if self.training else (torch.cat(z, 1), x)
替换为:
return x if self.training else torch.cat(z, 1)
最后如图所示:
models/export.py
改写后,官方的export.py已不适用,使用以下export代码:
'''
export y