Hey,
I want to use quantization aware training to quantize my model to int8. Unfortunately, I cant simply quantize the entire model, since my first layer is a batch normalization (after InputLayer), so I need to use a custom quantizationConfig for that layer. My problem is that I have a accuracy drop of around 4%, while using post-quantization the drop is only 2%. Is there anything wrong in the following code?
class DefaultDenseQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):

def get_weights_and_quantizers(self, layer):
  return [(layer.weights[i], LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False)) for i in range(2)]

def get_activations_and_quantizers(self, layer):
  return []

def set_quantize_weights(self, layer, quantize_weights):
  layer.weights[0] = quantize_weights[0]
  layer.weights[1] = quantize_weights[1]

def set_quantize_activations(self, layer, quantize_activations):
  pass

def get_output_quantizers(self, layer):
    return [tfmot.quantization.keras.quantizers.MovingAverageQuantizer(
        num_bits=8, per_axis=False, symmetric=False, narrow_range=False)]

def get_config(self):
  return {}

def apply_quantization_to_dense(self,layer):
    if layer != self.base_model.layers[1]:
        return tfmot.quantization.keras.quantize_annotate_layer(layer)
    else:
        return tfmot.quantization.keras.quantize_annotate_layer(
            layer,
            quantize_config=DefaultDenseQuantizeConfig())

 annotated_model = tf.keras.models.clone_model(
            self.base_model,
            clone_function=self.apply_quantization_to_dense,
        )
q_aware_model = quantize_annotate_model(annotated_model)
with quantize_scope(
    {'DefaultDenseQuantizeConfig': DefaultDenseQuantizeConfig,}):
        q_aware_model = tfmot.quantization.keras.quantize_apply(q_aware_model)

 q_aware_model.compile(optimizer=self.optimizer, loss=self.loss, loss_weights=self.loss_weights)
 q_aware_model.fit(self.train_dataset,
                          steps_per_epoch= int(self.ds_train_size / self.batch_size),     #
                          epochs=self.quantization_dict["quantization_epochs"],
                          callbacks=callbacks, )

converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
quantized_tflite_model = converter.convert()