glenn-jocher
commited on
Commit
•
8125ec5
1
Parent(s):
5402753
TFLite `--int8` 'flatbuffers==1.12' fix (#6216)
Browse files* TFLite `--int8` 'flatbuffers==1.12' fix
Temporary workaround for TFLite INT8 export.
* Update export.py
* Update export.py
export.py
CHANGED
@@ -277,8 +277,6 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
|
|
277 |
try:
|
278 |
import tensorflow as tf
|
279 |
|
280 |
-
from models.tf import representative_dataset_gen
|
281 |
-
|
282 |
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
|
283 |
batch_size, ch, *imgsz = list(im.shape) # BCHW
|
284 |
f = str(file).replace('.pt', '-fp16.tflite')
|
@@ -288,6 +286,8 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
|
|
288 |
converter.target_spec.supported_types = [tf.float16]
|
289 |
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
290 |
if int8:
|
|
|
|
|
291 |
dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
|
292 |
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
|
293 |
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
|
|
|
277 |
try:
|
278 |
import tensorflow as tf
|
279 |
|
|
|
|
|
280 |
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
|
281 |
batch_size, ch, *imgsz = list(im.shape) # BCHW
|
282 |
f = str(file).replace('.pt', '-fp16.tflite')
|
|
|
286 |
converter.target_spec.supported_types = [tf.float16]
|
287 |
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
288 |
if int8:
|
289 |
+
from models.tf import representative_dataset_gen
|
290 |
+
check_requirements(('flatbuffers==1.12',)) # https://github.com/ultralytics/yolov5/issues/5707
|
291 |
dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
|
292 |
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
|
293 |
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
|