Merge pull request #189 from ziqi-jin/master

fix export bugs

export.py

@@ -23,7 +23,7 @@
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', type=str, default=r'C:\Users\chen\Desktop\Model_Zoo\model_zoo/v5lite-e.pt', help='weights path')  # from yolov5/models/
     parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size')  # height, width
-    parser.add_argument('--concat', type=str, default=True, help='concat or not')
+    parser.add_argument('--concat', action='store_false', help='concat or not')
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
     parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')

models/yolo.py

@@ -48,13 +48,20 @@ def forward(self, x):
             x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
 
             if not self.training:  # inference
-                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
+                if torch.onnx.is_in_onnx_export():
+                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
+                elif self.grid[i].shape[2:4] != x[i].shape[2:4]:
                     self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
                 logits = x[i][..., 5:]
 
                 y = x[i].sigmoid()
-                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
+                if not torch.onnx.is_in_onnx_export():
+                    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
+                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i].data  # wh
+                    y = torch.cat((xy, wh, y[..., 4:]), -1)
                 z.append(y.view(bs, -1, self.no))
                 logits_.append(logits.view(bs, -1, self.no - 5))