New features and fixes

utils/export_damoyolo.py

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         boxes = x[1]
         scores, classes = torch.max(x[0], 2, keepdim=True)
-        return torch.cat((boxes, scores, classes.float()), dim=2)
+        return boxes, scores, classes
 
 
 def suppress_warnings():
@@ -65,21 +65,27 @@ def main(args):
 
     img_size = args.size * 2 if len(args.size) == 1 else args.size
 
-    onnx_input_im = torch.zeros(1, 3, *img_size).to(device)
+    onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
     onnx_output_file = cfg.miscs['exp_name'] + '.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'output': {
+        'boxes': {
+            0: 'batch'
+        },
+        'scores': {
+            0: 'batch'
+        },
+        'classes': {
             0: 'batch'
         }
     }
 
     print('Exporting the model to ONNX')
     torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['output'],
+                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
                       dynamic_axes=dynamic_axes if args.dynamic else None)
 
     if args.simplify:
@@ -100,11 +106,14 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=11, help='ONNX opset version')
     parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
     parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
+    parser.add_argument('--batch', type=int, default=1, help='Implicit batch-size')
     args = parser.parse_args()
     if not os.path.isfile(args.weights):
         raise SystemExit('Invalid weights file')
     if not os.path.isfile(args.config):
         raise SystemExit('Invalid config file')
+    if args.dynamic and args.batch > 1:
+        raise SystemExit('Cannot set dynamic batch-size and implicit batch-size at same time')
     return args