I am getting the following error when trying to do quantization aware training with tensorflow 2.7:
ValueError: `to_quantize` can only either be a tf.keras Sequential or Functional model.
The error occurs when calling this method:
quantize_model = tfmot.quantization.keras.quantize_model(model)
The model is defined below. I suppose the reason is that subclassed models are not supported? I have already trained(normal training, not QAT) multiple models with the definition below. Post-training quantization works, but i would like to try quantization aware training to see if it improves performance. Is there a way to be able to do quantization aware training with the model below, or alternatively define it in another way and redo normal training.
import tensorflow as tf from tensorflow import keras from tensorflow.keras.layers import Dense, Flatten, Activation, Dropout from tensorflow.keras.models import Model class Mobilenet_v2_transfer(Model): def __init__(self): super(Mobilenet_v2_transfer, self).__init__() self.base = tf.keras.applications.mobilenet_v2.MobileNetV2( input_shape=(224, 224, 3), alpha=1.0, include_top=False, weights='imagenet', pooling='avg') self.base.trainable = True for layer in self.base.layers[:130]: layer.trainable = False self.flatten = Flatten() self.dense = Dense(1, kernel_regularizer=tf.keras.regularizers.L2(0.01) self.sigmoid = Activation('sigmoid') def call(self, x): x = self.base(x) x = self.flatten(x) x = self.dense(x) x = self.sigmoid(x) return x