diff --git a/export.py b/export.py index 7feb525711e8..600e0c318f33 100644 --- a/export.py +++ b/export.py @@ -8,6 +8,7 @@ TorchScript | yolov5s.torchscript | `torchscript` ONNX | yolov5s.onnx | `onnx` CoreML | yolov5s.mlmodel | `coreml` +OpenVINO | yolov5s_openvino_model/ | `openvino` TensorFlow SavedModel | yolov5s_saved_model/ | `saved_model` TensorFlow GraphDef | yolov5s.pb | `pb` TensorFlow Lite | yolov5s.tflite | `tflite` @@ -15,13 +16,14 @@ TensorRT | yolov5s.engine | `engine` Usage: - $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs + $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml openvino saved_model tflite tfjs Inference: $ python path/to/detect.py --weights yolov5s.pt yolov5s.torchscript yolov5s.onnx yolov5s.mlmodel (under development) + yolov5s_openvino_model (under development) yolov5s_saved_model yolov5s.pb yolov5s.tflite @@ -144,6 +146,23 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')): return ct_model +def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')): + # YOLOv5 OpenVINO export + try: + check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/ + import openvino.inference_engine as ie + + LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') + f = str(file).replace('.pt', '_openvino_model' + os.sep) + + cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}" + subprocess.check_output(cmd, shell=True) + + LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + except Exception as e: + LOGGER.info(f'\n{prefix} export failure: {e}') + + 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, conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')): @@ -317,7 +336,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' imgsz=(640, 640), # image (height, width) batch_size=1, # batch size device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu - include=('torchscript', 'onnx', 'coreml'), # include formats + include=('torchscript', 'onnx'), # include formats half=False, # FP16 half-precision export inplace=False, # set YOLOv5 Detect() inplace=True train=False, # model.train() mode @@ -325,7 +344,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' int8=False, # CoreML/TF INT8 quantization dynamic=False, # ONNX/TF: dynamic axes simplify=False, # ONNX: simplify model - opset=14, # ONNX: opset version + opset=12, # ONNX: opset version verbose=False, # TensorRT: verbose log workspace=4, # TensorRT: workspace size (GB) nms=False, # TF: add NMS to model @@ -338,9 +357,12 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' t = time.time() include = [x.lower() for x in include] tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports - imgsz *= 2 if len(imgsz) == 1 else 1 # expand file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights) + # Checks + imgsz *= 2 if len(imgsz) == 1 else 1 # expand + opset = 12 if ('openvino' in include) else opset # OpenVINO requires opset <= 12 + # Load PyTorch model device = select_device(device) 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' # Exports if 'torchscript' in include: export_torchscript(model, im, file, optimize) - if 'onnx' in include: + if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX export_onnx(model, im, file, opset, train, dynamic, simplify) if 'engine' in include: export_engine(model, im, file, train, half, simplify, workspace, verbose) if 'coreml' in include: export_coreml(model, im, file) + if 'openvino' in include: + export_openvino(model, im, file) # TensorFlow Exports if any(tf_exports): @@ -413,7 +437,7 @@ def parse_opt(): parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') - parser.add_argument('--opset', type=int, default=14, help='ONNX: opset version') + parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version') parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log') parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)') parser.add_argument('--nms', action='store_true', help='TF: add NMS to model') diff --git a/requirements.txt b/requirements.txt index 22b51fc490e3..96fc9d1a1f32 100755 --- a/requirements.txt +++ b/requirements.txt @@ -27,6 +27,7 @@ seaborn>=0.11.0 # scikit-learn==0.19.2 # CoreML quantization # tensorflow>=2.4.1 # TFLite export # tensorflowjs>=3.9.0 # TF.js export +# openvino-dev # OpenVINO export # Extras -------------------------------------- # albumentations>=1.0.3