Yolov5 TensorRT推理加速(c++版)

Yolov5 不做赘述，目前目标检测里使用非常多的模型，效果和速度兼顾，性能强悍，配合TensorRT推理加速，在工业界可以说是非常流行的组合。

废话不多说，直接开整，以下使用的Tensor RT部署推理路线为：Pytorch-> ONNX -> TensorRT。

pytorch导出到onnx模型，可以非常方便，并且支持dynamic维度，配合netron工具，可以查看模型的网络结构，而TensorRT对ONNX的支持也非常完整，所以选择这一套流程，可以非常轻松的完成TensorRT的部署，同时，tensor RT提供官方的nms插件，使得推理代码可以免去编写nms的部分，极大提高效率。

GitHub仓库地址： gentlemanarch/yolov5-tensorrt

环境准备

1. 系统：Ubuntu20.04LTS系统，或者tensorRT官方docker镜像：nvcr.io/nvidia/tensorrt:21.05-py3（推荐）

2. TensorRT: 本篇使用的TensorRT7.2.3

3. Yolov5: 截止2021.09.11的develop分支代码

4. gcc: 9.3.0

5. torch: 1.8.2

6. onnx:1.10.1

7. 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 yolov5 .pt model to onnx model

Usage:
    python models/export.py --weigh