glenn-jocher commited on
Commit
00d7b97
1 Parent(s): 8125ec5

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

Browse files

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

Reorganizes #6216 fix to update before `tensorflow` import so no restart required.

* Update export.py

Files changed (1) hide show
  1. export.py +2 -1
export.py CHANGED
@@ -287,7 +287,6 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
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]
@@ -435,6 +434,8 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
435
  # TensorFlow Exports
436
  if any(tf_exports):
437
  pb, tflite, edgetpu, tfjs = tf_exports[1:]
 
 
438
  assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
439
  model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
440
  agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, topk_all=topk_all,
 
287
  converter.optimizations = [tf.lite.Optimize.DEFAULT]
288
  if int8:
289
  from models.tf import representative_dataset_gen
 
290
  dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data
291
  converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
292
  converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
 
434
  # TensorFlow Exports
435
  if any(tf_exports):
436
  pb, tflite, edgetpu, tfjs = tf_exports[1:]
437
+ if (tflite or edgetpu) and int8: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
438
+ check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
439
  assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
440
  model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
441
  agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, topk_all=topk_all,