I’m Trying to test Quantization Aware Trainin from TensorFlow Lite. The following source code creates an AI model (variable: model) trained with the MNIST dataset (just 1 epoch for testing purpose). I saved after the model in a file (model.h5). When I’m creating a Quantization Aware model named q_aware_model directly from the model variable with the quantize_model method, it works.
But when I’m loading my model stored in the model.h5 file and trying a Quantization Aware model named q_aware_model1, it fails with error:
q_aware_model1 = quantize_model1(model1)
ValueError: to_quantize can only either be a keras Sequential or Functional model.

I don’t know why. what is the problem?
Thanks.
Eddy33

My code:

import numpy as np
import tensorflow as tf
from tensorflow import keras
from keras.datasets import mnist

import tf_keras as keras

(x_train, y_train), (x_test, y_test) = mnist.load_data()


x_train = x_train.reshape(60000, 784)
x_test = x_test.reshape(10000, 784)
x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255
print(x_train.shape[0], 'train samples')
print(x_test.shape[0], 'test samples')

num_classes = 10

# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

from keras.models import Sequential
from keras import models, layers
from keras import regularizers

model = keras.Sequential([
keras.layers.Dropout(0.2,input_shape=(784,)),
keras.layers.Dense(1000, kernel_regularizer = regularizers.l2(0.01), activation='relu'),
keras.layers.Dropout(0.5),
keras.layers.Dense(1000, kernel_regularizer = regularizers.l2(0.01), activation='relu'),
keras.layers.Dropout(0.5),
keras.layers.Dense(10,  activation='softmax')
])   

#model.summary()

model.compile(loss=keras.losses.categorical_crossentropy, 
          optimizer='adam', 
          metrics=['accuracy'])

hist = model.fit(x_train, y_train,
                    batch_size=128,
                    epochs=1, # just 1
                    verbose=1,
                    validation_data=(x_test,y_test))

score = model.evaluate(x_test, y_test, verbose=1)
print("Test loss {:.4f}, accuracy {:.2f}%".format(score[0], score[1] * 100))

print("Saved model.h5 to disk")
model.save("model.h5")


# Quantization Aware Training
import tensorflow_model_optimization as tfmot

print("nnnDirect QAT")
quantize_model = tfmot.quantization.keras.quantize_model

q_aware_model = quantize_model(model)

q_aware_model.compile(loss=keras.losses.categorical_crossentropy, 
          optimizer='adam', 
          metrics=['accuracy'])

q_aware_model.summary()


print("nnnQAT from loading model.h5")
model1 = tf.keras.models.load_model('model.h5')
quantize_model1 = tfmot.quantization.keras.quantize_model

q_aware_model1 = quantize_model1(model1)

q_aware_model1.compile(loss=keras.losses.categorical_crossentropy, 
          optimizer='adam', 
          metrics=['accuracy'])

q_aware_model1.summary()