Commit
•
80473a6
1
Parent(s):
f3085ac
Update `export.py` with Detect, Validate usages (#6280)
Browse files
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 |
-
|
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 |
-
|
452 |
-
|
|
|
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
|
|
|
|
|
|
|
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='
|
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
|