mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

TFLite `--int8` 'flatbuffers==1.12' fix (#6216) 

* TFLite `--int8` 'flatbuffers==1.12' fix

Temporary workaround for TFLite INT8 export.

* Update export.py

* Update export.py
pull/6217/head
Glenn Jocher 2022-01-05 13:01:21 -08:00 committed by GitHub
parent 5402753a53
commit 8125ec5d42
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
export.py
View File

@ -277,8 +277,6 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
try: try:
import tensorflow as tf import tensorflow as tf
from models.tf import representative_dataset_gen
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
batch_size, ch, *imgsz = list(im.shape) # BCHW batch_size, ch, *imgsz = list(im.shape) # BCHW
f = str(file).replace('.pt', '-fp16.tflite') f = str(file).replace('.pt', '-fp16.tflite')
@ -288,6 +286,8 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
converter.target_spec.supported_types = [tf.float16] converter.target_spec.supported_types = [tf.float16]
converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.optimizations = [tf.lite.Optimize.DEFAULT]
if int8: if int8:
from models.tf import representative_dataset_gen
check_requirements(('flatbuffers==1.12',)) # https://github.com/ultralytics/yolov5/issues/5707
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]