From 12130a8c03fb1710b55a5572a9bbb1dec2ae5a61 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 10 May 2021 12:13:44 +0200
Subject: [PATCH] Add --optimize argument (#3093)

Fix for c++ runtime errors in https://github.com/ultralytics/yolov5/issues/2973
---
 models/export.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/models/export.py b/models/export.py
index acf83622aa75..3415088c787a 100644
--- a/models/export.py
+++ b/models/export.py
@@ -30,6 +30,7 @@
     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')
+    parser.add_argument('--optimize', action='store_true', help='optimize TorchScript for mobile')  # TorchScript-only
     parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')  # ONNX-only
     parser.add_argument('--simplify', action='store_true', help='simplify ONNX model')  # ONNX-only
     opt = parser.parse_args()
@@ -78,7 +79,7 @@
         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).save(f)  # https://pytorch.org/tutorials/recipes/script_optimized.html
+        (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}')
@@ -123,7 +124,6 @@
         import coremltools as ct
 
         print(f'{prefix} starting export with coremltools {ct.__version__}...')
-        # convert model from torchscript and apply pixel scaling as per detect.py
         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)