Failed to apply the QAT function 'quantize_model' to the sequential model that is defined using tensorflow.keras

[wwwind]

Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is here. An issue for anything not supported should be a feature request.

Describe the bug A clear and concise description of what the bug is.

System information

TensorFlow version (installed from source or binary): 2.17.0

TensorFlow Model Optimization version (installed from source or binary): 0.8.0

Python version: 3.10

Describe the expected behavior

Describe the current behavior Failed with the error

  File "/home/jamesbond/work/venv/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py", line 135, in quantize_model
    raise ValueError(
ValueError: `to_quantize` can only either be a keras Sequential or Functional model.

Code to reproduce the issue

import tensorflow as tf
print(tf.__version__)

import tensorflow_model_optimization as tfmot
print(tfmot.__version__)

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense

def create_model():

# Define a simple Sequential model
    model = Sequential([
    Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    MaxPooling2D((2, 2)),
    Flatten(),
    Dense(128, activation='relu'),
    Dense(10, activation='softmax')
    ])

    return model

to_quantize_model = create_model()
print(f'Type of the model is {type(to_quantize_model)}')

# Check that the model is sequential
if not isinstance(to_quantize_model, Sequential):
    raise ValueError('not sequantial')

# Check the whole condition from https://github.com/tensorflow/model-optimization/blob/ed3f0176b561fe693a3cc55b53a3605b943b6bbf/tensorflow_model_optimization/python/core/quantization/keras/quantize.py#L135
if not isinstance(to_quantize_model, Sequential) and not (
      hasattr(to_quantize_model, '_is_graph_network')
      and to_quantize_model._is_graph_network
  ):  # pylint: disable=protected-access
    raise ValueError(
        'Condition FAILED: `to_quantize` can only either be a keras Sequential or '
        'Functional model.'
    )

# But now this API fails with the condition above
qat_model = tfmot.quantization.keras.quantize_model(to_quantize_model)

Screenshots If applicable, add screenshots to help explain your problem.

Additional context My model is defined using tensorflow.keras instead of tensorflow_model_optimization.python.core.keras.compat like in tutorials and this leads to this error as model is not recognized as Sequential, although it is Sequential.

[hbellafkir]

I had the same issue. I just switched to tensorflow_model_optimization.python.core.keras.compat. The code didn’t require much modification.

[Litschi123]

TFMOT is using Keras version 2 which is only used by default until TensorFlow version 2.15.

1. Install tf_keras pip install tf_keras
2. Add os.environ['TF_USE_LEGACY_KERAS'] = "1" before importing TensorFlow in your script to make it use Keras v2.

For more details check: https://keras.io/getting_started/#tensorflow--keras-2-backwards-compatibility