Unverified 提交 80473a65 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

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

上级 f3085acc
...@@ -82,6 +82,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:' ...@@ -82,6 +82,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
ts.save(str(f), _extra_files=extra_files) ts.save(str(f), _extra_files=extra_files)
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return f
except Exception as e: except Exception as e:
LOGGER.info(f'{prefix} export failure: {e}') LOGGER.info(f'{prefix} export failure: {e}')
...@@ -125,7 +126,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst ...@@ -125,7 +126,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
except Exception as e: except Exception as e:
LOGGER.info(f'{prefix} simplifier failure: {e}') LOGGER.info(f'{prefix} simplifier failure: {e}')
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'") return f
except Exception as e: except Exception as e:
LOGGER.info(f'{prefix} export failure: {e}') LOGGER.info(f'{prefix} export failure: {e}')
...@@ -143,13 +144,13 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')): ...@@ -143,13 +144,13 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
subprocess.check_output(cmd, shell=True) subprocess.check_output(cmd, shell=True)
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
def export_coreml(model, im, file, prefix=colorstr('CoreML:')): def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
# YOLOv5 CoreML export # YOLOv5 CoreML export
ct_model = None
try: try:
check_requirements(('coremltools',)) check_requirements(('coremltools',))
import coremltools as ct import coremltools as ct
...@@ -162,10 +163,10 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')): ...@@ -162,10 +163,10 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
ct_model.save(f) ct_model.save(f)
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return ct_model, f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
return None, None
return ct_model
def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')): 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,7 +217,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
with builder.build_engine(network, config) as engine, open(f, 'wb') as t: with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
t.write(engine.serialize()) t.write(engine.serialize())
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
...@@ -225,7 +226,6 @@ def export_saved_model(model, im, file, dynamic, ...@@ -225,7 +226,6 @@ def export_saved_model(model, im, file, dynamic,
tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
conf_thres=0.25, prefix=colorstr('TensorFlow SavedModel:')): conf_thres=0.25, prefix=colorstr('TensorFlow SavedModel:')):
# YOLOv5 TensorFlow SavedModel export # YOLOv5 TensorFlow SavedModel export
keras_model = None
try: try:
import tensorflow as tf import tensorflow as tf
from tensorflow import keras from tensorflow import keras
...@@ -247,10 +247,10 @@ def export_saved_model(model, im, file, dynamic, ...@@ -247,10 +247,10 @@ def export_saved_model(model, im, file, dynamic,
keras_model.save(f, save_format='tf') keras_model.save(f, save_format='tf')
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return keras_model, f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
return None, None
return keras_model
def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): 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,6 +269,7 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False) tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
...@@ -300,7 +301,7 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te ...@@ -300,7 +301,7 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
tflite_model = converter.convert() tflite_model = converter.convert()
open(f, "wb").write(tflite_model) open(f, "wb").write(tflite_model)
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
...@@ -328,6 +329,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): ...@@ -328,6 +329,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
subprocess.run(cmd, shell=True, check=True) subprocess.run(cmd, shell=True, check=True)
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
...@@ -364,6 +366,7 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): ...@@ -364,6 +366,7 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
j.write(subst) j.write(subst)
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
return f
except Exception as e: except Exception as e:
LOGGER.info(f'\n{prefix} export failure: {e}') LOGGER.info(f'\n{prefix} export failure: {e}')
...@@ -431,15 +434,15 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' ...@@ -431,15 +434,15 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
# Exports # Exports
if 'torchscript' in include: if 'torchscript' in include:
export_torchscript(model, im, file, optimize) f = export_torchscript(model, im, file, optimize)
if 'engine' in include: # TensorRT required before ONNX if 'engine' in include: # TensorRT required before ONNX
export_engine(model, im, file, train, half, simplify, workspace, verbose) f = export_engine(model, im, file, train, half, simplify, workspace, verbose)
if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
export_onnx(model, im, file, opset, train, dynamic, simplify) f = export_onnx(model, im, file, opset, train, dynamic, simplify)
if 'openvino' in include: if 'openvino' in include:
export_openvino(model, im, file) f = export_openvino(model, im, file)
if 'coreml' in include: if 'coreml' in include:
export_coreml(model, im, file) _, f = export_coreml(model, im, file)
# TensorFlow Exports # TensorFlow Exports
if any(tf_exports): if any(tf_exports):
...@@ -447,22 +450,26 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' ...@@ -447,22 +450,26 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707 if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow` check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.' assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs, model, f = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, topk_all=topk_all, agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
conf_thres=conf_thres, iou_thres=iou_thres) # keras model topk_all=topk_all,
conf_thres=conf_thres, iou_thres=iou_thres) # keras model
if pb or tfjs: # pb prerequisite to tfjs if pb or tfjs: # pb prerequisite to tfjs
export_pb(model, im, file) f = export_pb(model, im, file)
if tflite or edgetpu: if tflite or edgetpu:
export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100) f = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
if edgetpu: if edgetpu:
export_edgetpu(model, im, file) f = export_edgetpu(model, im, file)
if tfjs: if tfjs:
export_tfjs(model, im, file) f = export_tfjs(model, im, file)
# Finish # Finish
LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)' LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
f"\nResults saved to {colorstr('bold', file.parent.resolve())}" f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
f'\nVisualize with https://netron.app') f"\nVisualize with https://netron.app"
f"\nDetect with `python detect.py --weights {f}`"
f" or `model = torch.hub.load('ultralytics/yolov5', 'custom', '{f}')"
f"\nValidate with `python val.py --weights {f}`")
def parse_opt(): def parse_opt():
...@@ -490,7 +497,7 @@ def parse_opt(): ...@@ -490,7 +497,7 @@ def parse_opt():
parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold') parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
parser.add_argument('--include', nargs='+', parser.add_argument('--include', nargs='+',
default=['torchscript', 'onnx'], default=['torchscript', 'onnx'],
help='available formats are (torchscript, onnx, engine, coreml, saved_model, pb, tflite, tfjs)') help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
opt = parser.parse_args() opt = parser.parse_args()
print_args(FILE.stem, opt) print_args(FILE.stem, opt)
return opt return opt
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论