OpenVINO Export (#6057)
Browse files* OpenVINO export
* Remove timeout
* Add 3 files
* str
* Constrain opset to 12
* Default ONNX opset to 12
* Make dir
* Make dir
* Cleanup
* Cleanup
* check_requirements(('openvino-dev',))
- export.py +30 -6
- requirements.txt +1 -0
export.py
CHANGED
@@ -8,6 +8,7 @@ PyTorch | yolov5s.pt | -
|
|
8 |
TorchScript | yolov5s.torchscript | `torchscript`
|
9 |
ONNX | yolov5s.onnx | `onnx`
|
10 |
CoreML | yolov5s.mlmodel | `coreml`
|
|
|
11 |
TensorFlow SavedModel | yolov5s_saved_model/ | `saved_model`
|
12 |
TensorFlow GraphDef | yolov5s.pb | `pb`
|
13 |
TensorFlow Lite | yolov5s.tflite | `tflite`
|
@@ -15,13 +16,14 @@ TensorFlow.js | yolov5s_web_model/ | `tfjs`
|
|
15 |
TensorRT | yolov5s.engine | `engine`
|
16 |
|
17 |
Usage:
|
18 |
-
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model
|
19 |
|
20 |
Inference:
|
21 |
$ python path/to/detect.py --weights yolov5s.pt
|
22 |
yolov5s.torchscript
|
23 |
yolov5s.onnx
|
24 |
yolov5s.mlmodel (under development)
|
|
|
25 |
yolov5s_saved_model
|
26 |
yolov5s.pb
|
27 |
yolov5s.tflite
|
@@ -144,6 +146,23 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
|
|
144 |
return ct_model
|
145 |
|
146 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
147 |
def export_saved_model(model, im, file, dynamic,
|
148 |
tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
|
149 |
conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')):
|
@@ -317,7 +336,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
317 |
imgsz=(640, 640), # image (height, width)
|
318 |
batch_size=1, # batch size
|
319 |
device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
320 |
-
include=('torchscript', 'onnx'
|
321 |
half=False, # FP16 half-precision export
|
322 |
inplace=False, # set YOLOv5 Detect() inplace=True
|
323 |
train=False, # model.train() mode
|
@@ -325,7 +344,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
325 |
int8=False, # CoreML/TF INT8 quantization
|
326 |
dynamic=False, # ONNX/TF: dynamic axes
|
327 |
simplify=False, # ONNX: simplify model
|
328 |
-
opset=
|
329 |
verbose=False, # TensorRT: verbose log
|
330 |
workspace=4, # TensorRT: workspace size (GB)
|
331 |
nms=False, # TF: add NMS to model
|
@@ -338,9 +357,12 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
338 |
t = time.time()
|
339 |
include = [x.lower() for x in include]
|
340 |
tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports
|
341 |
-
imgsz *= 2 if len(imgsz) == 1 else 1 # expand
|
342 |
file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)
|
343 |
|
|
|
|
|
|
|
|
|
344 |
# Load PyTorch model
|
345 |
device = select_device(device)
|
346 |
assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
|
@@ -372,12 +394,14 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
372 |
# Exports
|
373 |
if 'torchscript' in include:
|
374 |
export_torchscript(model, im, file, optimize)
|
375 |
-
if 'onnx' in include:
|
376 |
export_onnx(model, im, file, opset, train, dynamic, simplify)
|
377 |
if 'engine' in include:
|
378 |
export_engine(model, im, file, train, half, simplify, workspace, verbose)
|
379 |
if 'coreml' in include:
|
380 |
export_coreml(model, im, file)
|
|
|
|
|
381 |
|
382 |
# TensorFlow Exports
|
383 |
if any(tf_exports):
|
@@ -413,7 +437,7 @@ def parse_opt():
|
|
413 |
parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
|
414 |
parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
|
415 |
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
|
416 |
-
parser.add_argument('--opset', type=int, default=
|
417 |
parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
|
418 |
parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
|
419 |
parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
|
|
|
8 |
TorchScript | yolov5s.torchscript | `torchscript`
|
9 |
ONNX | yolov5s.onnx | `onnx`
|
10 |
CoreML | yolov5s.mlmodel | `coreml`
|
11 |
+
OpenVINO | yolov5s_openvino_model/ | `openvino`
|
12 |
TensorFlow SavedModel | yolov5s_saved_model/ | `saved_model`
|
13 |
TensorFlow GraphDef | yolov5s.pb | `pb`
|
14 |
TensorFlow Lite | yolov5s.tflite | `tflite`
|
|
|
16 |
TensorRT | yolov5s.engine | `engine`
|
17 |
|
18 |
Usage:
|
19 |
+
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml openvino saved_model tflite tfjs
|
20 |
|
21 |
Inference:
|
22 |
$ python path/to/detect.py --weights yolov5s.pt
|
23 |
yolov5s.torchscript
|
24 |
yolov5s.onnx
|
25 |
yolov5s.mlmodel (under development)
|
26 |
+
yolov5s_openvino_model (under development)
|
27 |
yolov5s_saved_model
|
28 |
yolov5s.pb
|
29 |
yolov5s.tflite
|
|
|
146 |
return ct_model
|
147 |
|
148 |
|
149 |
+
def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
|
150 |
+
# YOLOv5 OpenVINO export
|
151 |
+
try:
|
152 |
+
check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
|
153 |
+
import openvino.inference_engine as ie
|
154 |
+
|
155 |
+
LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
|
156 |
+
f = str(file).replace('.pt', '_openvino_model' + os.sep)
|
157 |
+
|
158 |
+
cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}"
|
159 |
+
subprocess.check_output(cmd, shell=True)
|
160 |
+
|
161 |
+
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
162 |
+
except Exception as e:
|
163 |
+
LOGGER.info(f'\n{prefix} export failure: {e}')
|
164 |
+
|
165 |
+
|
166 |
def export_saved_model(model, im, file, dynamic,
|
167 |
tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
|
168 |
conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')):
|
|
|
336 |
imgsz=(640, 640), # image (height, width)
|
337 |
batch_size=1, # batch size
|
338 |
device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
339 |
+
include=('torchscript', 'onnx'), # include formats
|
340 |
half=False, # FP16 half-precision export
|
341 |
inplace=False, # set YOLOv5 Detect() inplace=True
|
342 |
train=False, # model.train() mode
|
|
|
344 |
int8=False, # CoreML/TF INT8 quantization
|
345 |
dynamic=False, # ONNX/TF: dynamic axes
|
346 |
simplify=False, # ONNX: simplify model
|
347 |
+
opset=12, # ONNX: opset version
|
348 |
verbose=False, # TensorRT: verbose log
|
349 |
workspace=4, # TensorRT: workspace size (GB)
|
350 |
nms=False, # TF: add NMS to model
|
|
|
357 |
t = time.time()
|
358 |
include = [x.lower() for x in include]
|
359 |
tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports
|
|
|
360 |
file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)
|
361 |
|
362 |
+
# Checks
|
363 |
+
imgsz *= 2 if len(imgsz) == 1 else 1 # expand
|
364 |
+
opset = 12 if ('openvino' in include) else opset # OpenVINO requires opset <= 12
|
365 |
+
|
366 |
# Load PyTorch model
|
367 |
device = select_device(device)
|
368 |
assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
|
|
|
394 |
# Exports
|
395 |
if 'torchscript' in include:
|
396 |
export_torchscript(model, im, file, optimize)
|
397 |
+
if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
|
398 |
export_onnx(model, im, file, opset, train, dynamic, simplify)
|
399 |
if 'engine' in include:
|
400 |
export_engine(model, im, file, train, half, simplify, workspace, verbose)
|
401 |
if 'coreml' in include:
|
402 |
export_coreml(model, im, file)
|
403 |
+
if 'openvino' in include:
|
404 |
+
export_openvino(model, im, file)
|
405 |
|
406 |
# TensorFlow Exports
|
407 |
if any(tf_exports):
|
|
|
437 |
parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
|
438 |
parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
|
439 |
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
|
440 |
+
parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
|
441 |
parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
|
442 |
parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
|
443 |
parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
|
requirements.txt
CHANGED
@@ -27,6 +27,7 @@ seaborn>=0.11.0
|
|
27 |
# scikit-learn==0.19.2 # CoreML quantization
|
28 |
# tensorflow>=2.4.1 # TFLite export
|
29 |
# tensorflowjs>=3.9.0 # TF.js export
|
|
|
30 |
|
31 |
# Extras --------------------------------------
|
32 |
# albumentations>=1.0.3
|
|
|
27 |
# scikit-learn==0.19.2 # CoreML quantization
|
28 |
# tensorflow>=2.4.1 # TFLite export
|
29 |
# tensorflowjs>=3.9.0 # TF.js export
|
30 |
+
# openvino-dev # OpenVINO export
|
31 |
|
32 |
# Extras --------------------------------------
|
33 |
# albumentations>=1.0.3
|