Browse Source

Add `--int8` argument (#4799)

* Add `--int8` argument

* parents[0] bug fix

* Fix order
modifyDataloader
Glenn Jocher GitHub 3 years ago
parent
commit
b74dd4ba4f
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 5 deletions
  1. +8
    -5
      export.py

+ 8
- 5
export.py View File



FILE = Path(__file__).resolve() FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # yolov5/ dir ROOT = FILE.parents[0] # yolov5/ dir
sys.path.append(ROOT.as_posix()) # add yolov5/ to path
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH


from models.common import Conv from models.common import Conv
from models.experimental import attempt_load from models.experimental import attempt_load
print(f'\n{prefix} export failure: {e}') 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 # YOLOv5 TensorFlow Lite export
try: try:
import tensorflow as tf import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.optimizations = [tf.lite.Optimize.DEFAULT]
if tfl_int8:
if int8:
dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib) converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
inplace=False, # set YOLOv5 Detect() inplace=True inplace=False, # set YOLOv5 Detect() inplace=True
train=False, # model.train() mode train=False, # model.train() mode
optimize=False, # TorchScript: optimize for mobile 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 simplify=False, # ONNX: simplify model
opset=12, # ONNX: opset version opset=12, # ONNX: opset version
): ):
if pb or tfjs: # pb prerequisite to tfjs if pb or tfjs: # pb prerequisite to tfjs
export_pb(model, im, file) export_pb(model, im, file)
if tflite: 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: if tfjs:
export_tfjs(model, im, file) export_tfjs(model, im, file)


parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True') 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('--train', action='store_true', help='model.train() mode')
parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile') 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('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version') parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version')

Loading…
Cancel
Save