From 6dc4737d5620b9d5070d871aed7e1b85067f64c8 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 15 Sep 2021 11:25:28 +0200 Subject: [PATCH 1/3] Add `--int8` argument --- export.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/export.py b/export.py index 8d6805893d1e..b17cbb69fa60 100644 --- a/export.py +++ b/export.py @@ -32,8 +32,9 @@ from torch.utils.mobile_optimizer import optimize_for_mobile FILE = Path(__file__).resolve() -ROOT = FILE.parents[0] # yolov5/ dir -sys.path.append(ROOT.as_posix()) # add yolov5/ to path +ROOT = FILE.parents[1] # yolov5/ dir +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH from models.common import Conv from models.experimental import attempt_load @@ -174,7 +175,7 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): print(f'\n{prefix} export failure: {e}') -def export_tflite(keras_model, im, file, tfl_int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')): +def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')): # YOLOv5 TensorFlow Lite export try: import tensorflow as tf @@ -187,7 +188,7 @@ def export_tflite(keras_model, im, file, tfl_int8, data, ncalib, prefix=colorstr converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] converter.optimizations = [tf.lite.Optimize.DEFAULT] - if tfl_int8: + if int8: dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib) converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] @@ -233,6 +234,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' half=False, # FP16 half-precision export inplace=False, # set YOLOv5 Detect() inplace=True train=False, # model.train() mode + int8=False, # CoreML/TF INT8 quantization optimize=False, # TorchScript: optimize for mobile dynamic=False, # ONNX: dynamic axes simplify=False, # ONNX: simplify model @@ -288,7 +290,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' if pb or tfjs: # pb prerequisite to tfjs export_pb(model, im, file) if tflite: - export_tflite(model, im, file, tfl_int8=False, data=data, ncalib=100) + export_tflite(model, im, file, int8=int8, data=data, ncalib=100) if tfjs: export_tfjs(model, im, file) @@ -309,6 +311,7 @@ def parse_opt(): parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True') parser.add_argument('--train', action='store_true', help='model.train() mode') parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile') + 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=13, help='ONNX: opset version') From 6bf0d52660061eaf40515bd010c77b141c06c606 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 15 Sep 2021 11:26:45 +0200 Subject: [PATCH 2/3] parents[0] bug fix --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index b17cbb69fa60..16d3935994ac 100644 --- a/export.py +++ b/export.py @@ -32,7 +32,7 @@ from torch.utils.mobile_optimizer import optimize_for_mobile FILE = Path(__file__).resolve() -ROOT = FILE.parents[1] # yolov5/ dir +ROOT = FILE.parents[0] # yolov5/ dir if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATH From e9c29ba643eab6be548470ecae3dbc8bb351b17c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 15 Sep 2021 11:33:01 +0200 Subject: [PATCH 3/3] Fix order --- export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index 16d3935994ac..ea7f1ebd0b1f 100644 --- a/export.py +++ b/export.py @@ -234,9 +234,9 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' half=False, # FP16 half-precision export inplace=False, # set YOLOv5 Detect() inplace=True train=False, # model.train() mode - int8=False, # CoreML/TF INT8 quantization optimize=False, # TorchScript: optimize for mobile - dynamic=False, # ONNX: dynamic axes + int8=False, # CoreML/TF INT8 quantization + dynamic=False, # ONNX/TF: dynamic axes simplify=False, # ONNX: simplify model opset=12, # ONNX: opset version ):