Add grid concat and fuse such operators (#389)

* Add grid concat and fuse so many op

* Fix model

* Fix other detector

* Update yolo.py

* Update yolo.py

Co-authored-by: Alexey <AlexeyAB@users.noreply.github.com>

export.py

@@ -143,7 +143,8 @@ if __name__ == '__main__':
                     'output': {0: 'batch'},
                 }
             dynamic_axes.update(output_axes)
-        if opt.grid and opt.end2end:
+        if opt.grid:
+            if opt.end2end:
                 print('\nStarting export end2end onnx model for %s...' % 'TensorRT' if opt.max_wh is None else 'onnxruntime')
                 model = End2End(model,opt.topk_all,opt.iou_thres,opt.conf_thres,opt.max_wh,device)
                 if opt.end2end and opt.max_wh is None:
@@ -152,6 +153,8 @@ if __name__ == '__main__':
                               opt.batch_size, opt.topk_all, opt.batch_size, opt.topk_all]
                 else:
                     output_names = ['output']
+            else:
+                model.model[-1].concat = True
 
         torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
                           output_names=output_names,

models/yolo.py

@@ -25,6 +25,7 @@ class Detect(nn.Module):
     export = False  # onnx export
     end2end = False
     include_nms = False
+    concat = False
 
     def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
         super(Detect, self).__init__()
@@ -55,9 +56,10 @@ class Detect(nn.Module):
                     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:
-                    xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
+                    xy, wh, conf = y.split((2, 2, self.nc + 1), 4)  # y.tensor_split((2, 4, 5), 4)  # torch 1.8.0
-                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].data  # wh
+                    xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
-                    y = torch.cat((xy, wh, y[..., 4:]), -1)
+                    wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
+                    y = torch.cat((xy, wh, conf), 4)
                 z.append(y.view(bs, -1, self.no))
 
         if self.training:
@@ -67,6 +69,8 @@ class Detect(nn.Module):
         elif self.include_nms:
             z = self.convert(z)
             out = (z, )
+        elif self.concat:
+            out = torch.cat(z, 1)
         else:
             out = (torch.cat(z, 1), x)
 
@@ -95,6 +99,7 @@ class IDetect(nn.Module):
     export = False  # onnx export
     end2end = False
     include_nms = False
+    concat = False
 
     def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
         super(IDetect, self).__init__()
@@ -150,9 +155,10 @@ class IDetect(nn.Module):
                     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:
-                    xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
+                    xy, wh, conf = y.split((2, 2, self.nc + 1), 4)  # y.tensor_split((2, 4, 5), 4)  # torch 1.8.0
-                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].data  # wh
+                    xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
-                    y = torch.cat((xy, wh, y[..., 4:]), -1)
+                    wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
+                    y = torch.cat((xy, wh, conf), 4)
                 z.append(y.view(bs, -1, self.no))
 
         if self.training:
@@ -162,6 +168,8 @@ class IDetect(nn.Module):
         elif self.include_nms:
             z = self.convert(z)
             out = (z, )
+        elif self.concat:
+            out = torch.cat(z, 1)            
         else:
             out = (torch.cat(z, 1), x)
 
@@ -305,6 +313,7 @@ class IAuxDetect(nn.Module):
     export = False  # onnx export
     end2end = False
     include_nms = False
+    concat = False
 
     def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
         super(IAuxDetect, self).__init__()
@@ -344,9 +353,10 @@ class IAuxDetect(nn.Module):
                     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:
-                    xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
+                    xy, wh, conf = y.split((2, 2, self.nc + 1), 4)  # y.tensor_split((2, 4, 5), 4)  # torch 1.8.0
-                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].data  # wh
+                    xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5))  # new xy
-                    y = torch.cat((xy, wh, y[..., 4:]), -1)
+                    wh = wh ** 2 * (4 * self.anchor_grid[i].data)  # new wh
+                    y = torch.cat((xy, wh, conf), 4)
                 z.append(y.view(bs, -1, self.no))
 
         return x if self.training else (torch.cat(z, 1), x[:self.nl])
@@ -381,6 +391,8 @@ class IAuxDetect(nn.Module):
         elif self.include_nms:
             z = self.convert(z)
             out = (z, )
+        elif self.concat:
+            out = torch.cat(z, 1)            
         else:
             out = (torch.cat(z, 1), x)