glenn-jocher commited on
Commit
80473a6
1 Parent(s): f3085ac

Update `export.py` with Detect, Validate usages (#6280)

Browse files
Files changed (1) hide show
  1. export.py +30 -23
export.py CHANGED
@@ -82,6 +82,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
82
  ts.save(str(f), _extra_files=extra_files)
83
 
84
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
85
  except Exception as e:
86
  LOGGER.info(f'{prefix} export failure: {e}')
87
 
@@ -125,7 +126,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
125
  except Exception as e:
126
  LOGGER.info(f'{prefix} simplifier failure: {e}')
127
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
128
- LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
129
  except Exception as e:
130
  LOGGER.info(f'{prefix} export failure: {e}')
131
 
@@ -143,13 +144,13 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
143
  subprocess.check_output(cmd, shell=True)
144
 
145
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
146
  except Exception as e:
147
  LOGGER.info(f'\n{prefix} export failure: {e}')
148
 
149
 
150
  def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
151
  # YOLOv5 CoreML export
152
- ct_model = None
153
  try:
154
  check_requirements(('coremltools',))
155
  import coremltools as ct
@@ -162,10 +163,10 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
162
  ct_model.save(f)
163
 
164
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
165
  except Exception as e:
166
  LOGGER.info(f'\n{prefix} export failure: {e}')
167
-
168
- return ct_model
169
 
170
 
171
  def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
@@ -216,7 +217,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
216
  with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
217
  t.write(engine.serialize())
218
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
219
-
220
  except Exception as e:
221
  LOGGER.info(f'\n{prefix} export failure: {e}')
222
 
@@ -225,7 +226,6 @@ def export_saved_model(model, im, file, dynamic,
225
  tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
226
  conf_thres=0.25, prefix=colorstr('TensorFlow SavedModel:')):
227
  # YOLOv5 TensorFlow SavedModel export
228
- keras_model = None
229
  try:
230
  import tensorflow as tf
231
  from tensorflow import keras
@@ -247,10 +247,10 @@ def export_saved_model(model, im, file, dynamic,
247
  keras_model.save(f, save_format='tf')
248
 
249
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
250
  except Exception as e:
251
  LOGGER.info(f'\n{prefix} export failure: {e}')
252
-
253
- return keras_model
254
 
255
 
256
  def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
@@ -269,6 +269,7 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
269
  tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
270
 
271
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
272
  except Exception as e:
273
  LOGGER.info(f'\n{prefix} export failure: {e}')
274
 
@@ -300,7 +301,7 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
300
  tflite_model = converter.convert()
301
  open(f, "wb").write(tflite_model)
302
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
303
-
304
  except Exception as e:
305
  LOGGER.info(f'\n{prefix} export failure: {e}')
306
 
@@ -328,6 +329,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
328
  subprocess.run(cmd, shell=True, check=True)
329
 
330
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
331
  except Exception as e:
332
  LOGGER.info(f'\n{prefix} export failure: {e}')
333
 
@@ -364,6 +366,7 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
364
  j.write(subst)
365
 
366
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
367
  except Exception as e:
368
  LOGGER.info(f'\n{prefix} export failure: {e}')
369
 
@@ -431,15 +434,15 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
431
 
432
  # Exports
433
  if 'torchscript' in include:
434
- export_torchscript(model, im, file, optimize)
435
  if 'engine' in include: # TensorRT required before ONNX
436
- export_engine(model, im, file, train, half, simplify, workspace, verbose)
437
  if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
438
- export_onnx(model, im, file, opset, train, dynamic, simplify)
439
  if 'openvino' in include:
440
- export_openvino(model, im, file)
441
  if 'coreml' in include:
442
- export_coreml(model, im, file)
443
 
444
  # TensorFlow Exports
445
  if any(tf_exports):
@@ -447,22 +450,26 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
447
  if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
448
  check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
449
  assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
450
- model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
451
- agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, topk_all=topk_all,
452
- conf_thres=conf_thres, iou_thres=iou_thres) # keras model
 
453
  if pb or tfjs: # pb prerequisite to tfjs
454
- export_pb(model, im, file)
455
  if tflite or edgetpu:
456
- export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
457
  if edgetpu:
458
- export_edgetpu(model, im, file)
459
  if tfjs:
460
- export_tfjs(model, im, file)
461
 
462
  # Finish
463
  LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
464
  f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
465
- f'\nVisualize with https://netron.app')
 
 
 
466
 
467
 
468
  def parse_opt():
@@ -490,7 +497,7 @@ def parse_opt():
490
  parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
491
  parser.add_argument('--include', nargs='+',
492
  default=['torchscript', 'onnx'],
493
- help='available formats are (torchscript, onnx, engine, coreml, saved_model, pb, tflite, tfjs)')
494
  opt = parser.parse_args()
495
  print_args(FILE.stem, opt)
496
  return opt
 
82
  ts.save(str(f), _extra_files=extra_files)
83
 
84
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
85
+ return f
86
  except Exception as e:
87
  LOGGER.info(f'{prefix} export failure: {e}')
88
 
 
126
  except Exception as e:
127
  LOGGER.info(f'{prefix} simplifier failure: {e}')
128
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
129
+ return f
130
  except Exception as e:
131
  LOGGER.info(f'{prefix} export failure: {e}')
132
 
 
144
  subprocess.check_output(cmd, shell=True)
145
 
146
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
147
+ return f
148
  except Exception as e:
149
  LOGGER.info(f'\n{prefix} export failure: {e}')
150
 
151
 
152
  def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
153
  # YOLOv5 CoreML export
 
154
  try:
155
  check_requirements(('coremltools',))
156
  import coremltools as ct
 
163
  ct_model.save(f)
164
 
165
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
166
+ return ct_model, f
167
  except Exception as e:
168
  LOGGER.info(f'\n{prefix} export failure: {e}')
169
+ return None, None
 
170
 
171
 
172
  def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
 
217
  with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
218
  t.write(engine.serialize())
219
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
220
+ return f
221
  except Exception as e:
222
  LOGGER.info(f'\n{prefix} export failure: {e}')
223
 
 
226
  tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
227
  conf_thres=0.25, prefix=colorstr('TensorFlow SavedModel:')):
228
  # YOLOv5 TensorFlow SavedModel export
 
229
  try:
230
  import tensorflow as tf
231
  from tensorflow import keras
 
247
  keras_model.save(f, save_format='tf')
248
 
249
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
250
+ return keras_model, f
251
  except Exception as e:
252
  LOGGER.info(f'\n{prefix} export failure: {e}')
253
+ return None, None
 
254
 
255
 
256
  def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
 
269
  tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
270
 
271
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
272
+ return f
273
  except Exception as e:
274
  LOGGER.info(f'\n{prefix} export failure: {e}')
275
 
 
301
  tflite_model = converter.convert()
302
  open(f, "wb").write(tflite_model)
303
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
304
+ return f
305
  except Exception as e:
306
  LOGGER.info(f'\n{prefix} export failure: {e}')
307
 
 
329
  subprocess.run(cmd, shell=True, check=True)
330
 
331
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
332
+ return f
333
  except Exception as e:
334
  LOGGER.info(f'\n{prefix} export failure: {e}')
335
 
 
366
  j.write(subst)
367
 
368
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
369
+ return f
370
  except Exception as e:
371
  LOGGER.info(f'\n{prefix} export failure: {e}')
372
 
 
434
 
435
  # Exports
436
  if 'torchscript' in include:
437
+ f = export_torchscript(model, im, file, optimize)
438
  if 'engine' in include: # TensorRT required before ONNX
439
+ f = export_engine(model, im, file, train, half, simplify, workspace, verbose)
440
  if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
441
+ f = export_onnx(model, im, file, opset, train, dynamic, simplify)
442
  if 'openvino' in include:
443
+ f = export_openvino(model, im, file)
444
  if 'coreml' in include:
445
+ _, f = export_coreml(model, im, file)
446
 
447
  # TensorFlow Exports
448
  if any(tf_exports):
 
450
  if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
451
  check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
452
  assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
453
+ model, f = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
454
+ agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
455
+ topk_all=topk_all,
456
+ conf_thres=conf_thres, iou_thres=iou_thres) # keras model
457
  if pb or tfjs: # pb prerequisite to tfjs
458
+ f = export_pb(model, im, file)
459
  if tflite or edgetpu:
460
+ f = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
461
  if edgetpu:
462
+ f = export_edgetpu(model, im, file)
463
  if tfjs:
464
+ f = export_tfjs(model, im, file)
465
 
466
  # Finish
467
  LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
468
  f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
469
+ f"\nVisualize with https://netron.app"
470
+ f"\nDetect with `python detect.py --weights {f}`"
471
+ f" or `model = torch.hub.load('ultralytics/yolov5', 'custom', '{f}')"
472
+ f"\nValidate with `python val.py --weights {f}`")
473
 
474
 
475
  def parse_opt():
 
497
  parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
498
  parser.add_argument('--include', nargs='+',
499
  default=['torchscript', 'onnx'],
500
+ help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
501
  opt = parser.parse_args()
502
  print_args(FILE.stem, opt)
503
  return opt