Quantisation of custom model with custom layer (Full int 8)

user_9010507 0
Viewed 147

Hi I have a custom model, that requires every layer to be int8 quantised.

I have a custom layer called CustomLayer

Post training quantisation (Works):

converter = tf.lite.TFLiteConverter.from_keras_model(tflite_model)
converter.optimizations = [tf.lite.Optimize.q]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8  # or tf.uint8
converter.inference_output_type = tf.int8  # or tf.uint8
print("converting")
tflite_model_bytes = converter.convert()
print("converted")
model_file_name = 'tflite_model.tflite'

# Save the model.
with open(model_file_name, 'wb') as f:
    f.write(tflite_model_bytes)

Quantise aware training approach:

I am using default quantise config

class DefaultDenseQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):
    # Configure how to quantize weights.
    def get_weights_and_quantizers(self, layer):
      return [(layer.kernel, LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False))]

    # Configure how to quantize activations.
    def get_activations_and_quantizers(self, layer):
      return [(layer.activation, MovingAverageQuantizer(num_bits=8, symmetric=False, narrow_range=False, per_axis=False))]

    def set_quantize_weights(self, layer, quantize_weights):
      # Add this line for each item returned in `get_weights_and_quantizers`
      # , in the same order
      layer.kernel = quantize_weights[0]

    def set_quantize_activations(self, layer, quantize_activations):
      # Add this line for each item returned in `get_activations_and_quantizers`
      # , in the same order.
      layer.activation = quantize_activations[0]

    # Configure how to quantize outputs (may be equivalent to activations).
    def get_output_quantizers(self, layer):
      return []

    def get_config(self):
      return {}

And trying to convert using

with quantize_scope({
    "DefaultDenseQuantizeConfig": DefaultDenseQuantizeConfig,
    "CustomLayer": CustomLayer
}):

    q_aware_model = tfmot.quantization.keras.quantize_model(tflite_model)

I get this error:

RuntimeError: Layer custom_layer_1:<class 'CustomLayer'> is not supported. You can quantize this layer by passing a tfmot.quantization.keras.QuantizeConfig instance to the quantize_annotate_layer API.

So I try annotating my layers approach

Annotating layers:

I have tried annotating my layers within my model using this:

with quantize_scope({
    "DefaultDenseQuantizeConfig": DefaultDenseQuantizeConfig,
    "CustomLayer": CustomLayer
}):
    def apply_quantization_to_layer(layer):
      return tfmot.quantization.keras.quantize_annotate_layer(layer, DefaultDenseQuantizeConfig())
    
    annotated_model = tf.keras.models.clone_model(
        tflite_model,
        clone_function=apply_quantization_to_layer,
    )

    tfmot.quantization.keras.quantize_apply(annotated_model)

AttributeError: 'CustomLayer' object has no attribute 'kernel'

This error is occurring in

def get_weights_and_quantizers(self, layer):
      return [(layer.kernel, LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False))]

And has the same problem in get_activations_and_quantizers for layer.activation

tfmot keras tensorflow python
1 Answers

Ok I think I've found a solution as I am no longer receiving any errors, it seems to train and it seems to save as a tflite file successfully.

What I assume happens is when get_weights_and_quantizers returns a list each item is (tensorFromLayer, quantizerYouWantToUse) it then uses the quantizerYouWantToUse to quantise the tensorFromLayer. In set_quantize_weights it returns the list with the quantised tensors which you then set reset for the layer.

This is a solution I came up with:

class DefaultDenseQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):

    # List all of your weights
    weights = {
        "kernel": LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False)
    }

    # List of all your activations
    activations = {
        "activation": MovingAverageQuantizer(num_bits=8, symmetric=False, narrow_range=False, per_axis=False)
    }

    # Configure how to quantize weights.
    def get_weights_and_quantizers(self, layer):
        output = []
        for attribute, quantizer in self.weights.items():
            if hasattr(layer, attribute):
                output.append((getattr(layer, attribute), quantizer))

        return output

    # Configure how to quantize activations.
    def get_activations_and_quantizers(self, layer):
        output = []
        for attribute, quantizer in self.activations.items():
            if hasattr(layer, attribute):
                output.append((getattr(layer, attribute), quantizer))

        return output

    def set_quantize_weights(self, layer, quantize_weights):
        # Add this line for each item returned in `get_weights_and_quantizers`
        # , in the same order

        count = 0
        for attribute in self.weights.keys():
            if hasattr(layer, attribute):
                setattr(layer, attribute, quantize_weights[count])
                count += 1

    def set_quantize_activations(self, layer, quantize_activations):
        # Add this line for each item returned in `get_activations_and_quantizers`
        # , in the same order.
        count = 0
        for attribute in self.activations.keys():
            if hasattr(layer, attribute):
                setattr(layer, attribute, quantize_activations[count])
                count += 1

    # Configure how to quantize outputs (may be equivalent to activations).
    def get_output_quantizers(self, layer):
      return []

    def get_config(self):
      return {}
user_90105070
Related