カスタム レイヤーを使用した量子化 - 2.5 日本語

Vitis AI ユーザー ガイド (UG1414)

Document ID
UG1414
Release Date
2022-06-15
Version
2.5 日本語

Tensorflow 2 には機械学習モデルを構築するための一般的なビルトイン レイヤーが数多く用意されています。また、特定用途向けのレイヤーをゼロから作成することも、既存レイヤーを組み合わせて作成することも容易です。「レイヤー」は tf.keras の中心的なアブストラクションの 1 つで、カスタム レイヤーの作成には Layer をサブクラス化する方法が推奨されます。詳細は、TensorFlow ユーザー ガイド を参照してください。

Vai_q_tensorflow2 は、カスタム レイヤーを使用したモデルの量子化など、サブクラス化により新しいカスタム レイヤーをサポートします。また、カスタム量子化ストラテジを使用したカスタム レイヤーの量子化も実験的にサポートしています。

注記: 現在のリリースの vai_q_tensorflow2 では、tf.keras.Model のサブクラス化によるカスタム モデルはサポートされないため、レイヤーにフラット化してください。

カスタム レイヤーを含むモデルの量子化

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 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 によってラップされるため、量子化されません。完全なサンプルは、ここ にあります。

注記: カスタム レイヤーは量子化されないため、運用中に dump_model API を呼び出してデータ チェック用のゴールデン結果をダンプする際は、dump_float=True に設定し、カスタム レイヤーの浮動小数点の重みとアクティベーションをダンプする必要があります。

(実験的機能) カスタム量子化ストラテジを使用したカスタム レイヤーの量子化

カスタム レイヤーは、vai_q_tensorflow2 のサポートされる API リストに含まれないため、デフォルトの量子化ストラテジではカスタム レイヤーは量子化されず、量子化中も浮動小数点のまま残ります。カスタム量子化ストラテジを構築して量子化を実行できるように、上級ユーザー向けに 'custom_quantize_strategy' という名前のインターフェイスが用意されています。カスタム量子化ストラテジは、Dict の JSON ファイルとして量子化ストラテジを項目に持つ Dict オブジェクトです。

Default Quantize Strategy には、量子化ストラテジのサンプルが用意されています。カスタム量子化ストラテジも同じフォーマットに従います。同じ項目がある場合は、カスタム量子化ストラテジの項目によってデフォルト ストラテジの項目がオーバーライドされ、新しい項目は量子化ストラテジに追加されます。

この機能を使用すると、上記のサンプルに示した '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 quantization 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)