Add `--include torchscript onnx coreml` argument (#3137)

* Allow users to skip exporting in formats that they don't care about

* Correct comments

* Update export.py

renamed --skip-format to --exclude

* Switched format from exclude to include (as instructed by @glenn-jocher)

* cleanup

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

models/export.py

@@ -1,7 +1,7 @@
-"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats
+"""Exports a YOLOv5 *.pt model to TorchScript, ONNX, CoreML formats
 
 Usage:
-    $ export PYTHONPATH="$PWD" && python models/export.py --weights yolov5s.pt --img 640 --batch 1
+    $ python path/to/models/export.py --weights yolov5s.pt --img 640 --batch 1
 """
 
 import argparse
@@ -27,6 +27,7 @@ if __name__ == '__main__':
     parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats')
     parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
     parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
     parser.add_argument('--train', action='store_true', help='model.train() mode')
@@ -35,6 +36,7 @@ if __name__ == '__main__':
     parser.add_argument('--simplify', action='store_true', help='simplify ONNX model')  # ONNX-only
     opt = parser.parse_args()
     opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
+    opt.include = [x.lower() for x in opt.include]
     print(opt)
     set_logging()
     t = time.time()
@@ -47,7 +49,7 @@ if __name__ == '__main__':
     # Checks
     gs = int(max(model.stride))  # grid size (max stride)
     opt.img_size = [check_img_size(x, gs) for x in opt.img_size]  # verify img_size are gs-multiples
-    assert not (opt.device.lower() == "cpu" and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
+    assert not (opt.device.lower() == 'cpu' and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
 
     # Input
     img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device)  # image size(1,3,320,192) iDetection
@@ -74,62 +76,66 @@ if __name__ == '__main__':
     print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
 
     # TorchScript export -----------------------------------------------------------------------------------------------
-    prefix = colorstr('TorchScript:')
-    try:
-        print(f'\n{prefix} starting export with torch {torch.__version__}...')
-        f = opt.weights.replace('.pt', '.torchscript.pt')  # filename
-        ts = torch.jit.trace(model, img, strict=False)
-        (optimize_for_mobile(ts) if opt.optimize else ts).save(f)
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-    except Exception as e:
-        print(f'{prefix} export failure: {e}')
+    if 'torchscript' in opt.include or 'coreml' in opt.include:
+        prefix = colorstr('TorchScript:')
+        try:
+            print(f'\n{prefix} starting export with torch {torch.__version__}...')
+            f = opt.weights.replace('.pt', '.torchscript.pt')  # filename
+            ts = torch.jit.trace(model, img, strict=False)
+            (optimize_for_mobile(ts) if opt.optimize else ts).save(f)
+            print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        except Exception as e:
+            print(f'{prefix} export failure: {e}')
 
     # ONNX export ------------------------------------------------------------------------------------------------------
-    prefix = colorstr('ONNX:')
-    try:
-        import onnx
-
-        print(f'{prefix} starting export with onnx {onnx.__version__}...')
-        f = opt.weights.replace('.pt', '.onnx')  # filename
-        torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
-                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
-                                        'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
-
-        # Checks
-        model_onnx = onnx.load(f)  # load onnx model
-        onnx.checker.check_model(model_onnx)  # check onnx model
-        # print(onnx.helper.printable_graph(model_onnx.graph))  # print
-
-        # Simplify
-        if opt.simplify:
-            try:
-                check_requirements(['onnx-simplifier'])
-                import onnxsim
-
-                print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
-                model_onnx, check = onnxsim.simplify(model_onnx,
-                                                     dynamic_input_shape=opt.dynamic,
-                                                     input_shapes={'images': list(img.shape)} if opt.dynamic else None)
-                assert check, 'assert check failed'
-                onnx.save(model_onnx, f)
-            except Exception as e:
-                print(f'{prefix} simplifier failure: {e}')
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-    except Exception as e:
-        print(f'{prefix} export failure: {e}')
+    if 'onnx' in opt.include:
+        prefix = colorstr('ONNX:')
+        try:
+            import onnx
+
+            print(f'{prefix} starting export with onnx {onnx.__version__}...')
+            f = opt.weights.replace('.pt', '.onnx')  # filename
+            torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
+                              dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
+                                            'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
+
+            # Checks
+            model_onnx = onnx.load(f)  # load onnx model
+            onnx.checker.check_model(model_onnx)  # check onnx model
+            # print(onnx.helper.printable_graph(model_onnx.graph))  # print
+
+            # Simplify
+            if opt.simplify:
+                try:
+                    check_requirements(['onnx-simplifier'])
+                    import onnxsim
+
+                    print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
+                    model_onnx, check = onnxsim.simplify(
+                        model_onnx,
+                        dynamic_input_shape=opt.dynamic,
+                        input_shapes={'images': list(img.shape)} if opt.dynamic else None)
+                    assert check, 'assert check failed'
+                    onnx.save(model_onnx, f)
+                except Exception as e:
+                    print(f'{prefix} simplifier failure: {e}')
+            print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        except Exception as e:
+            print(f'{prefix} export failure: {e}')
 
     # CoreML export ----------------------------------------------------------------------------------------------------
-    prefix = colorstr('CoreML:')
-    try:
-        import coremltools as ct
-
-        print(f'{prefix} starting export with coremltools {ct.__version__}...')
-        model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
-        f = opt.weights.replace('.pt', '.mlmodel')  # filename
-        model.save(f)
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-    except Exception as e:
-        print(f'{prefix} export failure: {e}')
+    if 'coreml' in opt.include:
+        prefix = colorstr('CoreML:')
+        try:
+            import coremltools as ct
+
+            print(f'{prefix} starting export with coremltools {ct.__version__}...')
+            model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
+            f = opt.weights.replace('.pt', '.mlmodel')  # filename
+            model.save(f)
+            print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        except Exception as e:
+            print(f'{prefix} export failure: {e}')
 
     # Finish
     print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.')