TFLite Model used in official documentation doesn't compile on Edge TPU Compiler

[gaikwadrahul8]

1. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 11
• TensorFlow installation (pip package or built from source): pip
• TensorFlow library (version, if pip package or github SHA, if built from source):
  • pip show tensorflow = 2.17.0
  • tf.version = 2.18.0-dev20240815

2. Code

Using the code from Post-training integer quantization official tutorial to create and convert a TensorFlow model to TFlite:

import logging
logging.getLogger("tensorflow").setLevel(logging.DEBUG)

import tensorflow as tf
import numpy as np
print("TensorFlow version: ", tf.__version__)

# Load MNIST dataset
mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Normalize the input image so that each pixel value is between 0 to 1.
train_images = train_images.astype(np.float32) / 255.0
test_images = test_images.astype(np.float32) / 255.0

# Define the model architecture
model = tf.keras.Sequential([
  tf.keras.layers.InputLayer(shape=(28, 28)),
  tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
  tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
  tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(10)
])

# Train the digit classification model
model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(
                  from_logits=True),
              metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=5, validation_data=(test_images, test_labels))

def representative_data_gen():
  for input_value in tf.data.Dataset.from_tensor_slices(train_images).batch(1).take(100):
    yield [input_value]

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
# Ensure that if any ops can't be quantized, the converter throws an error
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.target_spec.supported_types = [tf.int8] 
# Set the input and output tensors to uint8 (APIs added in r2.3)
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.experimental_new_quantizer = True

tflite_model_quant = converter.convert()

import pathlib

tflite_models_dir = pathlib.Path("mnist_tflite_models/")
tflite_models_dir.mkdir(exist_ok=True, parents=True)
tflite_model_quant_file = tflite_models_dir/"mnist_model_quant.tflite"
tflite_model_quant_file.write_bytes(tflite_model_quant)

Which displays the following output:

2024-08-16 18:07:19.094940: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-08-16 18:07:20.845501: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
TensorFlow version:  2.18.0-dev20240815
2024-08-16 18:07:25.704749: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
Epoch 1/5
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 7s 3ms/step - accuracy: 0.8672 - loss: 0.4860 - val_accuracy: 0.9722 - val_loss: 0.0965
Epoch 2/5
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 7s 4ms/step - accuracy: 0.9738 - loss: 0.0919 - val_accuracy: 0.9768 - val_loss: 0.0735
Epoch 3/5
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 8s 4ms/step - accuracy: 0.9797 - loss: 0.0696 - val_accuracy: 0.9799 - val_loss: 0.0622
Epoch 4/5
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 7s 3ms/step - accuracy: 0.9833 - loss: 0.0565 - val_accuracy: 0.9800 - val_loss: 0.0600
Epoch 5/5
1875/1875 ━━━━━━━━━━━━━━━━━━━━ 6s 3ms/step - accuracy: 0.9843 - loss: 0.0524 - val_accuracy: 0.9813 - val_loss: 0.0595
INFO:tensorflow:Assets written to: C:\Users\Me\AppData\Local\Temp\tmptd9h6442\assets
INFO:tensorflow:Assets written to: C:\Users\Me\AppData\Local\Temp\tmptd9h6442\assets
Saved artifact at 'C:\Users\Me\AppData\Local\Temp\tmptd9h6442'. The following endpoints are available:

* Endpoint 'serve'
  args_0 (POSITIONAL_ONLY): TensorSpec(shape=(None, 28, 28), dtype=tf.float32, name='keras_tensor')
Output Type:
  TensorSpec(shape=(None, 10), dtype=tf.float32, name=None)
Captures:
  1897857613456: TensorSpec(shape=(), dtype=tf.resource, name=None)
  1897859106576: TensorSpec(shape=(), dtype=tf.resource, name=None)
  1897857613072: TensorSpec(shape=(), dtype=tf.resource, name=None)
  1897857612688: TensorSpec(shape=(), dtype=tf.resource, name=None)
C:\Users\Me\.pyenv\pyenv-win\versions\3.12.4\Lib\site-packages\tensorflow\lite\python\convert.py:983: UserWarning: Statistics for quantized inputs were expected, but not specified; continuing anyway.
  warnings.warn(
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1723828081.771633   14420 tf_tfl_flatbuffer_helpers.cc:359] Ignored output_format.
W0000 00:00:1723828081.772124   14420 tf_tfl_flatbuffer_helpers.cc:362] Ignored drop_control_dependency.
2024-08-16 18:08:01.772979: I tensorflow/cc/saved_model/reader.cc:83] Reading SavedModel from: C:\Users\Me\AppData\Local\Temp\tmptd9h6442
2024-08-16 18:08:01.774123: I tensorflow/cc/saved_model/reader.cc:52] Reading meta graph with tags { serve }
2024-08-16 18:08:01.774323: I tensorflow/cc/saved_model/reader.cc:147] Reading SavedModel debug info (if present) from: C:\Users\Me\AppData\Local\Temp\tmptd9h6442
I0000 00:00:1723828081.779428   14420 mlir_graph_optimization_pass.cc:401] MLIR V1 optimization pass is not enabled
2024-08-16 18:08:01.780783: I tensorflow/cc/saved_model/loader.cc:236] Restoring SavedModel bundle.
2024-08-16 18:08:01.819711: I tensorflow/cc/saved_model/loader.cc:220] Running initialization op on SavedModel bundle at path: C:\Users\Me\AppData\Local\Temp\tmptd9h6442
2024-08-16 18:08:01.830586: I tensorflow/cc/saved_model/loader.cc:462] SavedModel load for tags { serve }; Status: success: OK. Took 57614 microseconds.
2024-08-16 18:08:01.847290: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2024-08-16 18:08:02.073373: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence
fully_quantize: 0, inference_type: 6, input_inference_type: UINT8, output_inference_type: UINT8

3. Failure after conversion

When I then run the newly created TFLite file through the edgetpu_compiler (via Docker) it fails saying it still has dynamic-sized tensors:

docker run --rm -it -v .:/home/edgetpu edgetpu-compiler edgetpu_compiler mnist_model_quant.tflite
Edge TPU Compiler version 16.0.384591198
Started a compilation timeout timer of 180 seconds.
ERROR: Attempting to use a delegate that only supports static-sized tensors with a graph that has dynamic-sized tensors.
Compilation failed: Model failed in Tflite interpreter. Please ensure model can be loaded/run in Tflite interpreter.
Compilation child process completed within timeout period.
Compilation failed!

Any idea how I can fully convert the model to static-sized tensors. I tried the suggestion of using converter._experimental_new_quantizer but that didn't help.

[pkgoogle]

Currently blocked by #150

[gaikwadrahul8]

This issue originally reported by @ADarkDividedGem has been moved to this dedicated repository for ai-edge-torch to enhance issue tracking and prioritization. To ensure continuity, we have created this new issue on your behalf.

We appreciate your understanding and look forward to your continued involvement.