MobileNetV3 QAT TFLite Conversion Issue

Empieza reproduciendo la conversión de QAT a TFLite de MobileNetV3Large con TensorFlow 2.15.0 y TensorFlow Model Optimization 0.7.5, utilizando la refactorización proporcionada de hard_sigmoid y CustomQuantizeConfig. Compara las salidas de QAT en Keras y TFLite, y después lee los Issues relacionados #368 y #974 para consultar la investigación previa. Se considera terminado cuando el modelo convertido ya no muestra la diferencia de precisión indicada.

Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is here. An issue for anything not supported should be a feature request.

Describe the bug
A clear and concise description of what the bug is.

This is a similar issue to #368 but for MobileNetV3Large, where after Quantisation Aware Training, I see a large drop in accuracy in the QAT TFLite model compared to the corresponding QAT Keras Model. Minor Implementation details: I had to refactor the default MobileNetV3Large Keras Code to make it compatible with QAT in the Tensorflow Model Optimisation library by replacing the Add operations in its Hard Sigmoid function with Rescaling and using a Moving Average Output only Quantiser for the Multiply and Rescaling layers in the network. I train the network for more than 6-7 epochs with ~5100 batches in each epoch (each batch consisting of 10 samples) but I see no convergence between the Keras and TFlite models as was seen in #368. I see a number of people in #974 have raised the same issue but this has not been fixed yet. It's likely that is a kernel implementation bug similar to #368 so would be great if a fix could be developed for this. It might be helpful to note that I didn't face this issue in MobileNetV3Large minimalistic which makes me wonder that the issue might be in the Multiply Layers of the Squeeze-Excite and Hard Swish functions.

Would appreciate any help. Thanks!

System information

TensorFlow version (installed from source or binary): 2.15.0

TensorFlow Model Optimization version (installed from source or binary): 0.7.5

Python version: 3.10.12

Keras Version: 2.15.0

Describe the expected behavior
QAT Keras Model should generate identical outputs to the converted QAT TFlite Model

Describe the current behavior
Converted QAT TFlite Model has much lower accuracy than the corresponding QAT Keras Model

Code to reproduce the issue
Provide a reproducible code that is the bare minimum necessary to generate the
problem.

This is the refactored part of MobileNetV3:

def hard_sigmoid(x):
    return layers.Rescaling(1.0 / 6.0, offset=0.0)(
        layers.ReLU(6.0)(layers.Rescaling(1.0, offset=3.0)(x))

The Quantization Config I use for Multiply and Rescaling layers is this:

class CustomQuantizeConfig(quantize_config.QuantizeConfig):
    """QuantizeConfig which only quantizes layer outputs."""

    def get_weights_and_quantizers(self, layer):
        return []

    def get_activations_and_quantizers(self, layer):
        return []

    def set_quantize_weights(self, layer, quantize_weights):
        pass

    def set_quantize_activations(self, layer, quantize_activations):
        pass

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

    def get_config(self):
        return {}

I've tried using an AllValuesQuantizer as well since it was mentioned in #368 that MovingAverageQuantizer takes time to converge but that didn't seem to help either.

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
The use case for the network is Monocular Depth Estimation so there is a decoding network attached on top of the MobileNetV3 encoder.