Tensorflow Quantization Aware Training

Question

I want to quantize a DenseNet model. I am using Tensorflow 2.4.

import tensorflow_model_optimization as tfmot
model = tf.keras.applications.DenseNet121(include_top=True,weights=None,input_tensor=None,input_shape=None,pooling=None,classes=1000) 
quantize_model = tfmot.quantization.keras.quantize_model
model = quantize_model(model)

But I got the following message:

RuntimeError: Layer conv2_block1_0_bn:<class 'tensorflow.python.keras.layers.normalization_v2.BatchNormalization'> is not supported. You can quantize this layer by passing a tfmot.quantization.keras.QuantizeConfig instance to the quantize_annotate_layer API.

Is there a way how I can do this. I can not change the keras code.

Answer 1

In your case you need to quantize the layer BatchNormalization seperately.

If you see the below example code snippet from this Quantization TF Guide , DefaultDenseQuantizeConfig is used to handle this problem. Hope this guide helpy you solve this.

quantize_annotate_layer = tfmot.quantization.keras.quantize_annotate_layer
quantize_annotate_model = tfmot.quantization.keras.quantize_annotate_model
quantize_scope = tfmot.quantization.keras.quantize_scope

class CustomLayer(tf.keras.layers.Dense):
  pass

model = quantize_annotate_model(tf.keras.Sequential([
   quantize_annotate_layer(CustomLayer(20, input_shape=(20,)), DefaultDenseQuantizeConfig()),
   tf.keras.layers.Flatten()
]))

# `quantize_apply` requires mentioning `DefaultDenseQuantizeConfig` with `quantize_scope`
# as well as the custom Keras layer.
with quantize_scope(
  {'DefaultDenseQuantizeConfig': DefaultDenseQuantizeConfig,
   'CustomLayer': CustomLayer}):
  # Use `quantize_apply` to actually make the model quantization aware.
  quant_aware_model = tfmot.quantization.keras.quantize_apply(model)

quant_aware_model.summary()