glenn-jocher commited on
Commit
525f4f8
·
unverified ·
1 Parent(s): 57b0d3a

Add --optimize argument (#3093)

Browse files

Fix for c++ runtime errors in https://github.com/ultralytics/yolov5/issues/2973

Files changed (1) hide show
  1. models/export.py +2 -2
models/export.py CHANGED
@@ -30,6 +30,7 @@ if __name__ == '__main__':
30
  parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
31
  parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
32
  parser.add_argument('--train', action='store_true', help='model.train() mode')
 
33
  parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only
34
  parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only
35
  opt = parser.parse_args()
@@ -78,7 +79,7 @@ if __name__ == '__main__':
78
  print(f'\n{prefix} starting export with torch {torch.__version__}...')
79
  f = opt.weights.replace('.pt', '.torchscript.pt') # filename
80
  ts = torch.jit.trace(model, img, strict=False)
81
- optimize_for_mobile(ts).save(f) # https://pytorch.org/tutorials/recipes/script_optimized.html
82
  print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
83
  except Exception as e:
84
  print(f'{prefix} export failure: {e}')
@@ -123,7 +124,6 @@ if __name__ == '__main__':
123
  import coremltools as ct
124
 
125
  print(f'{prefix} starting export with coremltools {ct.__version__}...')
126
- # convert model from torchscript and apply pixel scaling as per detect.py
127
  model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
128
  f = opt.weights.replace('.pt', '.mlmodel') # filename
129
  model.save(f)
 
30
  parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
31
  parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
32
  parser.add_argument('--train', action='store_true', help='model.train() mode')
33
+ parser.add_argument('--optimize', action='store_true', help='optimize TorchScript for mobile') # TorchScript-only
34
  parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only
35
  parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only
36
  opt = parser.parse_args()
 
79
  print(f'\n{prefix} starting export with torch {torch.__version__}...')
80
  f = opt.weights.replace('.pt', '.torchscript.pt') # filename
81
  ts = torch.jit.trace(model, img, strict=False)
82
+ (optimize_for_mobile(ts) if opt.optimize else ts).save(f)
83
  print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
84
  except Exception as e:
85
  print(f'{prefix} export failure: {e}')
 
124
  import coremltools as ct
125
 
126
  print(f'{prefix} starting export with coremltools {ct.__version__}...')
 
127
  model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
128
  f = opt.weights.replace('.pt', '.mlmodel') # filename
129
  model.save(f)