Quantization aware training for a Transfer Learning MobileNet model
publioelon opened this issue · 5 comments
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: #377 (comment)
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?
Any updates on this?
Any update on this ?
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)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.