カスタム レイヤーを使用した量子化 - 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 は、{"custom_layer_class_name":"custom_layer_class"} を含む dict で、複数のカスタム レイヤーはカンマで区切られます。また、カスタム レイヤーを使用してモデルを量子化する際は、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 インターフェイスを使用してカスタムの量子化ストラテジを作成して、量子化の実験ができます。

カスタム量子化ストラテジは、JSON 形式の量子化ストラテジ項目を持つ Dict オブジェクトとして表されます。

デフォルトの量子化ストラテジは、量子化ストラテジ形式の例を提供し、カスタムの量子化戦略は同じ構造に従います。同じ項目がある場合、カスタム量子化ストラテジの項目によってデフォルト ストラテジの項目がオーバーライドされ、新しい項目は量子化ストラテジに追加されます。

この機能を使用すると、カスタムの量子化ストラテジを使用して上記のサンプルに示した 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)