Quantizing with Custom Layers - 3.5 English

Vitis AI User Guide (UG1414)

Document ID
Release Date
3.5 English

TensorFlow 2 offers a rich set of built-in layers to construct machine learning models, and it also provides straightforward methods to create application-specific layers from scratch or by combining existing ones. The layer is a central abstraction in tf.keras, and subclassing this class is the recommended approach to developing custom layers. For more detailed information, refer to the TensorFlow user guide.

vai_q_tensorflow2 supports new custom layers through subclassing, which includes the capability to quantize models with custom layers. Also, it provides experimental support for quantizing these custom layers using custom quantize strategies.

Note: Custom model through subclassing tf.keras.T is not supported by vai_q_tensorflow2 in this release. Flatten it into layers.

Quantizing models with custom layers

vai_q_tensorflow2 provides interfaces to load the custom layers that are available in some models. For example:

class MyCustomLayer(keras.layers.Layer):

    def __init__(self, units=32, **kwargs):
        super(MyLayer, self).__init__(kwargs)
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
        self.b = self.add_weight(
            shape=(self.units,), initializer="zeros", trainable=True, name='b')

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

    def get_config(self):
        base_config = super(MyLayer, self).get_config()
        config = {"units": self.units}
        return dict(list(base_config.items()) + list(config.items()

# Here is a float model with custom layer," "MyCustomLayer", use the custom_objects argument in tf.keras.models.load_model to load it.
float_model = tf.keras.models.load_model('float_model.h5', custom_objects={'MyCustomLayer': MyCustomLayer})

The float model contains a custom layer named "MyCustomLayer" and the custom_objects argument in the tf.keras.model.load_model API loads it. Similarly, the VitisQuantizer class provides the 'custom_objects' argument to handle the custom layers. The following code is an example. The argument custom_objects is a dict containing the {"custom_layer_class_name":"custom_layer_class"}, and commas separate multiple custom layers. Moreover, add_shape_info should also be set to True for the quantize_model API when quantizing models with custom layers to add shape inference information for them.

from tensorflow_model_optimization.quantization.keras import vitis_quantize
# Register the custom layer to VitisQuantizer by custom_objects argument.
quantizer = vitis_quantize.VitisQuantizer(float_model, custom_objects={'MyCustomLayer': MyCustomLayer})
quantized_model = quantizer.quantize_model(calib_dataset=calib_dataset, calib_step=100, calib_batch_size=10, add_shape_info=True)

During the quantization, these custom layers are wrapped by CustomLayerWrapper and kept unquantized. For a complete example, click here.

Note: When using the dump_model API to generate golden results for data checking during deployment, set dump_float=True to dump float weights and activations for the custom layers, as these layers are not quantized.

(Experimental) Quantizing custom layers with custom quantize strategy

The default quantize strategy does not quantize custom layers, as they are not included in the list of supported APIs for vai_q_tensorflow2. However, advanced users can create custom quantize strategies using the custom_quantize_strategy interface to conduct quantization experiments.

The custom quantize strategy is represented as a Dict object containing the quantize strategy items in JSON format.

The default quantize strategy provides an example of the quantize strategy format, and the custom quantize strategy follows the same structure. However, any item in the custom quantize strategy overrides the corresponding one in the default strategy while new items are added to the quantize strategy.

With this feature, you can quantize the MyCustomLayer layer from the previous example using a custom quantize strategy.

# Define quantizer with custom quantize strategy, which quantizes w,b and outputs 0 of MyCustomLayer objects.
my_quantize_strategy = {
    "quantize_registry_config": {
        "layer_quantize_config": [{
            "layer_type": "__main__.MyCustomLayer",
            "quantizable_weights": ["w", "b"],
            "weight_quantizers": [
                "quantizer_type": "LastValueQuantPosQuantizer","quantizer_params": {"bit_width": 8, "method": 1, "round_mode": 0}, 
                "quantizer_type": "LastValueQuantPosQuantizer", "quantizer_params": {"bit_width": 8, "method": 1, "round_mode": 0}
            "quantizable_outputs": ["0"],
            "output_quantizers": [
                "quantizer_type": "LastValueQuantPosQuantizer", "quantizer_params": {"bit_width": 8, "method": 1, "round_mode": 1}
quantizer = vitis_quantize.VitisQuantizer(model, custom_objects={'MyLayer': MyLayer}, custom_quantize_strategy=my_quantize_strategy)

# The following quantizison process are all the same as before, here we do normal PTQ as an example
quantized_model = quantizer.quantize_model(calib_dataset=calib_dataset, calib_step=100, calib_batch_size=10)