利用定制层进行量化 - 3.5 简体中文

量化具有定制层的模型

某些模型具有定制层，vai_q_tensorflow2 提供了加载定制层的接口。例如：

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),
            initializer="random_normal",
            trainable=True,
            name='w')
        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})

浮点模型包含一个名为 "MyCustomLayer" 的定制层，tf.keras.model.load_model API 中的 custom_objects 实参用于加载该定制层。同样，VitisQuantizer 类提供 'custom_objects' 实参来处理定制层。下面给出了一个代码示例。custom_objects 实参是一个 dict，其中包含 {"custom_layer_class_name":"custom_layer_class"}，并以逗号来分隔多个定制层。此外，在量化含有定制层的模型时，quantize_model API 的 add_shape_info 也应设为 True，以便为这些定制层添加形状推断信息。

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)

在量化期间，这些定制层会通过 CustomLayerWrapper 进行封装并保持未量化状态。要获取完整示例，请单击此处。

（实验性）利用定制量化策略量化定制层

默认量化策略不会对定制层进行量化，因为它们不包含在 vai_q_tensorflow2 支持的 API 列表中。但高级用户可以使用 custom_quantize_strategy 接口创建定制量化策略以开展量化实验。

定制量化策略呈现为一个 Dict 对象，其中包含 JSON 格式的量化策略项。

默认量化策略提供了量化策略格式示例，定制量化策略也遵循相同结构。但定制量化策略中的任何项都会覆盖默认策略中的对应项，而新的项则会添加到量化策略中。

您可利用此功能来量化来自上述使用定制量化策略的示例中的 MyCustomLayer 层。

# 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)