【目标检测-YOLO】YOLOv5-yolov5s TensorRT部署准备之ONNX导出（第一篇）

下面的内容是对 YOLOv5-yolov5s TensorRT部署前的准备之导出可用的ONNX。

之前已经写过部分内容，但是还不够详细。

1. YOLOv5-6.0版本

以下内容以 6.0 版本为准。

修改导出文件：

    def forward(self, x):z = []  # inference outputfor i in range(self.nl):x[i] = self.m[i](x[i])  # conv2d# 修改1：对于任何用到shape、size返回值的参数时，为了避免pytorch 导出 onnx 时候，对size 进行跟踪，跟踪时候会生成gather、shape的节点。使用 map 加上int转换。bs, _, ny, nx = map(int,x[i].shape)  # x(bs,255,20,20) to x(bs,3,20,20,85)# 修改2：对于reshape、view操作时，-1的指定请放到batch维度。其他维度可以计算出来即可。batch维度禁止指定为大于-1的明确数字。所以这里 bs 改为 -1.x[i] = x[i].view(-1, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()# export.py：model.train() if train else model.eval()  # 264行： train=False,  # model.train() mode 所以这里是 model.eval()模式, self.training 为 False,所以推理使用以下代码if not self.training:  # inferenceif self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)y = x[i].sigmoid()# inplace 操作是直接在 y 上进行操作（5.0 版本中使用该操作）if self.inplace:  y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xyy[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh# 非 inplace 操作，新建 xy, wh 变量，cat后赋值给 yelse:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953# (1, 3(na), 80, 80, 2) + (1, 3(na), 80, 80, 2) * 8xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy# (1, 3(na), 80, 80, 2) * (1, 3(na), 80, 80, 2)wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # why = torch.cat((xy, wh, y[..., 4:]), -1)  # -1 最后一个维度拼接 bs,3,80,80,85# 修改3：原因同 修改2。bs为-1，int(y.size(1)*y.size(2)*y.size(3)) 为 3*80*80或3*40*40或3*20*20z.append(y.view(-1, int(y.size(1)*y.size(2)*y.size(3)), self.no)) # -1,3*80*80,85  no = 85# 修改4：去掉下面的x, 节省内存，最终模型输出一个节点return x if self.training else torch.cat(z, 1)# grid 网格 如 3*3 网格def _make_grid(self, nx=20, ny=20, i=0):d = self.anchors[i].device  # 设备 cpu 或 'cuda:...'# yv = [0, 0, 0  #       1, 1, 1#       2, 2, 2]# xv = [0, 1, 2#       0, 1, 2#       0, 1, 2]yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)])# nx,ny,2 -> 1, na, nx,ny,2 # 参考：https://www.cnblogs.com/yanghailin/p/14329117.htmlgrid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()# 把 anchors * stride 得到 anchor_grid 还原到原图大小 一个head：shape: na(3),2 -> 1, na, 1, 1, 2 -> 1, na, ny,nx,2anchor_grid = (self.anchors[i].clone() * self.stride[i]) \.view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()return grid, anchor_grid

导出 onnx 文件：为了看模型各节点输出尺寸，我们加上 --simplify

python export.py --include onnx --simplify

YOLOv5有3个head，假设输入图像大小为640，那么输出为：

Bs *128*80*80；Bs*256*40*40；Bs*512*20*20；

如下图中Detect所示：

注：这里onnx图中，Bs = 1。

Bs*128*80*80； Bs*256*40*40； Bs*512*20*20； ------》》》

Bs*3*80*80*85; Bs*3*40*40*85; Bs*3*20*20*85;

注：3为每个cell 的Anchors数目，85 = x, y, w, h, obj, 80类（排序就是这样）。

推理模式：

export.py:264: train=False, # model.train() mode

295：model.train() if train else model.eval() # training mode = no Detect() layer grid construction # 这里是 model.eval()

注：这里 Stride 为模型下采样的倍数，YOLOv5 3个head 下采样为：8(80)；16(40)；32(20)。

上图中的 C_x与C_y 对应代码中的 self.grid；p_w*Stride和p_h*Stride对应代码中的 anchor_grid。

上图中的操作:

代码中是 for 循环三个head输出。下面以 80*80 head 为例说明：

问题：这里 w, h 为什么没有乘 Stride ? 原因：anchor_grid 已经是 640 尺度下的，故 anchor_grid大小隐含了 Stride(上面代码中: self.anchors[i].clone() * self.stride[i])。

所以，最后每个head输出：

上图中的操作放到 onnx 中做，在C++ 推理decode该张量的时候就会很简单。后面C++代码中只需要 for 循环 25200次，每次取出满足条件的；然后再 NMS，就结束了。

补充：避免使用inplace操作，例如y[…, 0:2] = y[…, 0:2] * 2 - 0.5

代码中：

                # 训练时候没影响，导出onnx问题很大，v6.0 默认使用的 非inplace 操作if self.inplace:y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xyy[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh

导出时候如果：export.py

m.inplace = inplace

parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')

python export.py --include onnx --inplace

上图可知，多了很多节点，并且还有 ScatterND节点，这种需要插件支持，这么复杂推理时候容易失败，所以不能 --inplace 导出。

目前未理解问题：yolo.py 中的 if self.grid[i].shape[2:4] != x[i].shape[2:4] : 是啥意思？

未解决：我尝试debug 找到 grid 在哪里进行的初始化，但是我并没有找到，只找到了：

models/experimental.py: model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model ->

models/yolo.py: m = self.model[-1] # Detect()

然后被初始化为:

x : [1, 255, 4, 4]，[1, 256, 2, 2]，[1, 512, 1, 1]

self.grid: [1, 1, 84, 44, 2]，[1, 1, 42, 22, 2]，[1, 1, 21, 11, 2]

84, 44 ！= 4, 4，就画网格。

真正推理时候：x : [1, 3, 80, 80, 85]， [1, 3, 40, 40, 85]， [1, 3, 20, 20, 85]

84, 44 != 80, 80；42, 22 != 40,40；21, 11 != 20,20；

然后就画网格呗。

得到： self.grid[i]：1, na, ny, nx, 2；self.anchor_grid[i] : 1, na, ny, nx, 2

        # nx,ny,2 -> 1, na, ny,nx,2grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float()# 把 anchors * stride 得到 anchor_grid 还原到原图大小 一个head：shape: na(3),2 -> 1, na, 1, 1, 2 -> 1, na, ny, nx, 2anchor_grid = (self.anchors[i].clone() * self.stride[i]) \.view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float()

上面代码的维度变换的目的是为了能够和输出维度对应起来便于操作。

2. YOLOv5-5.0版本

YOLOv5-5.0版本的默认onnx 导出没有进行上面的操作，也就是执行 python models/export.py 会得到3个输出。想得到合并后的输出可以使用：

python models/export.py --grid

原因：

        self.training |= self.export

    parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')model.model[-1].export = not opt.grid  # set Detect() layer grid export

然后按 6.0v 修改4处class Detect。

再使用上面命令导出，出现了 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# z.append(y.view(bs, -1, self.no))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)  # why = torch.cat((xy, wh, y[..., 4:]), -1)z.append(y.view(-1, int(y.size(1)*y.size(2)*y.size(3)), self.no))return x if self.training else torch.cat(z, 1)

注： self.anchor_grid[i].view(1, self.na, 1, 1, 2) 做了维度变换才能和前面的相乘。

参考：

详解TensorRT的C++/Python高性能部署，实战应用到项目_哔哩哔哩_bilibili