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