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

Files changed (1) hide show
  1. export.py +2 -2
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]