Quantization aware training for a Transfer Learning MobileNet model

[publioelon]

Hello, I have a MobileNetV2 That I am trying to use for image classification by means of transfer learning, although apparently seems to not work. Initially, I perform transfer learning on my model as follows:

base_model = tf.keras.applications.MobileNetV2(include_top=False)
base_model.trainable = True

inputs = tf.keras.layers.Input(shape=(384, 288, 3), name="input_layer")
x = tf.keras.layers.experimental.preprocessing.Rescaling(1./255)(inputs)
x = base_model(inputs)
x = tf.keras.layers.GlobalAveragePooling2D(name="global_average_pooling_layer")(x)
outputs = tf.keras.layers.Dense(5, activation="softmax", name="output_layer")(x)

model = tf.keras.Model(inputs, outputs)
model.compile(loss='categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

history = model.fit(training_dataset,
                                 epochs=30,
                                 steps_per_epoch=len(training_dataset),
                                 validation_data=validation_data,
                                 validation_steps=int(len(validation_data)))

I then followed the steps here to perform model quantization: https://colab.research.google.com/github/tensorflow/model-optimization/blob/master/tensorflow_model_optimization/g3doc/guide/quantization/training_example.ipynb#scrollTo=oq6blGjgFDCW

As the above, I attempted a quantization aware training on my model like this:

model = tf.keras.Model(inputs, outputs)

quantize_model = tfmot.quantization.keras.quantize_model

model = quantize_model(model)
model.compile(loss='categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

history = model.fit(training_dataset,
                                 epochs=30,
                                 steps_per_epoch=len(training_dataset),
                                 validation_data=validation_data,
                                 validation_steps=int(len(validation_data)))

It gives me the following error: ValueError: Quantizing a tf.keras Model inside another tf.keras Model is not supported.

Then I tried what was in this link: https://github.com/tensorflow/model-optimization/issues/377#issuecomment-820948555

I tried doing the below attempt:

base_model = tf.keras.applications.MobileNetV2(include_top=False)
base_model.trainable = False
q_base_model = quantize_model(base_model)

inputs = tf.keras.layers.Input(shape=(384, 288, 3), name="input_layer")
x = tf.keras.layers.experimental.preprocessing.Rescaling(1./255)(inputs)
x = base_model(inputs)
x = tf.keras.layers.GlobalAveragePooling2D(name="global_average_pooling_layer")(x)
outputs = tf.keras.layers.Dense(5, activation="softmax", name="output_layer")(x)

model = tf.keras.Model(inputs, outputs)
model = quantize_model(q_base_model)
model.compile(loss='categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

history = model.fit(training_dataset,
                                 epochs=30,
                                 steps_per_epoch=len(training_dataset),
                                 validation_data=validation_data,
                                 validation_steps=int(len(validation_data)))

Which outputs the following error:

ValueError                                Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py](https://localhost:8080/#) in quantize_apply(model, scheme)
    437   try:
--> 438     model_copy = _clone_model_with_weights(model)
    439   except ValueError:

15 frames
ValueError: Unknown layer: QuantizeLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py](https://localhost:8080/#) in quantize_apply(model, scheme)
    439   except ValueError:
    440     raise ValueError(
--> 441         'Unable to clone model. This generally happens if you used custom '
    442         'Keras layers or objects in your model. Please specify them via '
    443         '`quantize_scope` for your calls to `quantize_model` and '

ValueError: Unable to clone model. This generally happens if you used custom Keras layers or objects in your model. Please specify them via quantize_scope for your calls to quantize_model and quantize_apply. It seems I haven't fully understood how to get quantization aware training done correctly. I'd like to request help on how to properly do QAT on a transfer learning model such as the above example?

[thaink]

@Xhark could you take a look at this issue?

[gcunhase]

Any updates on this?

[Dinuka-Thathsara]

Any update on this ?

[mhyeonsoo]

Can you guys try to apply quantize to dense layer only?

I successfully QATed the model with the code below.

def apply_quantization_to_dense(layer):
        if isinstance(layer, tf.keras.layers.Dense):
            return tfmot.quantization.keras.quantize_annotate_layer(layer)
        return layer

    annotated_model = tf.keras.models.clone_model(
        model,
        clone_function=apply_quantization_to_dense,
    )

quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)

[Diaa340]

I found a workaround for this problem as I had the same issue. Basically, After reading the transfer learning model e.g.

import tensorflow as tf 
model_pre = tf.keras.applications.MobileNetV2(
    include_top=False,
    input_shape=(299, 299, 3),
    pooling='avg',
    weights='imagenet'
)
for layer in model_pre.layers:
    layer.trainable = False

you need to convert it into a keras model through

inputs = model_pre.input
outputs = model_pre.output
concatenated_model = tf.keras.Model(inputs=inputs, outputs=outputs)

then you can use the method @mhyeonsoo used to add quantization aware labels to the TF model. Note that using it directly without the previous approaches will not quantize the layers in the transfer learning model.

def apply_quantization_to_dense(layer):
        if isinstance(layer, tf.keras.layers.Dense):
            return tfmot.quantization.keras.quantize_annotate_layer(layer)
        return layer

    annotated_model = tf.keras.models.clone_model(
        model,
        clone_function=apply_quantization_to_dense,
    )

quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)

You can verify that your model is quantization aware from the model summary.