python - TensorFlow 的 QAT 似乎没有使用 AllValuesQuantizer 执行每通道量化

Question

我使用 AllValuesQuantizer 创建了两个 QAT 模型，一个使用每个张量，一个使用每个通道量化。在检查它们各自的 QuantizeWrapper 层时，我注意到它们都有变量 kernel_min 和 kernel_max 的标量值。

这是一个每张量量化模型的示例

这是每通道量化模型的示例

正如我从这篇论文中了解到的，内核的最小/最大值是定义比例和零点量化参数的东西。对于每张量量化，模型只有一个最小值和最大值是合理的，因为整个张量具有相同的尺度和零点。但是，对于每通道量化（每个通道都有自己的比例和零点）我相信 kernel_min 和 kernel_max 应该是向量？他们为什么不呢？

在这个 github 问题中，有人提到 QAT 自动使用每张量量化（截至 2020 年 3 月），但这可能会发生变化。对我来说，看起来 QAT 仍然只使用每张量量化？如果是这种情况，为什么我可以设置一个参数来启用每张量量化（参见 AllValuesQuantizer 的每轴布尔值）？

为了进一步说明我的观点，我还在AllValuesQuantizer 的源代码中指出 self.per_axis 永远不会传递给下一个函数，那么这个 even 变量是做什么用的？请注意，其他量化器 LastValue 和 MovingAverage 确实传递了此变量。

所以; TF 的 QAT 甚至会执行每通道量化吗？在我看来不像。如何通过 AllValuesQuantizer 使用每通道量化？

GitHub问题：https ://github.com/tensorflow/tensorflow/issues/47858

复制我的两个模型的代码：

import tensorflow as tf
from tensorflow import keras
import tensorflow_model_optimization as tfmot

from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession

config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)

# Possible quantization aware quantizers:
QAT_ALL_VALUES = tfmot.quantization.keras.quantizers.AllValuesQuantizer
QAT_LAST_VALUE = tfmot.quantization.keras.quantizers.LastValueQuantizer
QAT_MA = tfmot.quantization.keras.quantizers.MovingAverageQuantizer


def quantization_aware_training(model, save, w_bits, a_bits, symmetric, per_axis, narrow_range, quantizer, batch_size=64, epochs=2):

    # Create quantized model's name string
    name = model.name + '_'
    name = name + str(w_bits) + 'wbits_' + str(a_bits) + 'abits_'

    if symmetric:
        name = name + 'sym_'
    else:
        name = name + 'asym_'

    if narrow_range:
        name = name + 'narr_'
    else:
        name = name + 'full_'

    if per_axis:
        name = name + 'perch_'
    else:
        name = name + 'perten_'

    if quantizer == QAT_ALL_VALUES:
        name = name + 'AV'
    elif quantizer == QAT_LAST_VALUE:
        name = name + 'LV'
    elif quantizer == QAT_MA:
        name = name + 'MA'

    # Quantization
    # *****
    quantize_apply = tfmot.quantization.keras.quantize_apply
    quantize_model = tfmot.quantization.keras.quantize_model
    quantize_annotate_layer = tfmot.quantization.keras.quantize_annotate_layer
    clone_model = tf.keras.models.clone_model
    quantize_scope = tfmot.quantization.keras.quantize_scope

    supported_layers = [
        tf.keras.layers.Conv2D,
    ]

    class Quantizer(tfmot.quantization.keras.QuantizeConfig):
        # Configure how to quantize weights.
        def get_weights_and_quantizers(self, layer):
            return [(layer.kernel, tfmot.quantization.keras.quantizers.LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False))]

        # Configure how to quantize activations.
        def get_activations_and_quantizers(self, layer):
            return [(layer.activation, tfmot.quantization.keras.quantizers.MovingAverageQuantizer(num_bits=8, symmetric=False, narrow_range=False, per_axis=False))]

        def set_quantize_weights(self, layer, quantize_weights):
            # Add this line for each item returned in `get_weights_and_quantizers`
            # , in the same order
            layer.kernel = quantize_weights[0]

        def set_quantize_activations(self, layer, quantize_activations):
            # Add this line for each item returned in `get_activations_and_quantizers`
            # , in the same order.
            layer.activation = quantize_activations[0]

        # Configure how to quantize outputs (may be equivalent to activations).
        def get_output_quantizers(self, layer):
            return []

        def get_config(self):
            return {}

    class ConvQuantizer(Quantizer):
        # Configure weights to quantize with 4-bit instead of 8-bits.
        def get_weights_and_quantizers(self, layer):
            return [(layer.kernel, quantizer(num_bits=w_bits, symmetric=symmetric, narrow_range=narrow_range, per_axis=per_axis))]

        # Configure how to quantize activations.
        def get_activations_and_quantizers(self, layer):
            return [(layer.activation, tfmot.quantization.keras.quantizers.MovingAverageQuantizer(num_bits=a_bits, symmetric=False, narrow_range=False, per_axis=False))]

    class DepthwiseQuantizer(Quantizer):
        # Configure weights to quantize with 4-bit instead of 8-bits.
        def get_weights_and_quantizers(self, layer):
            return [(layer.depthwise_kernel, quantizer(num_bits=w_bits, symmetric=symmetric, narrow_range=narrow_range, per_axis=per_axis))]

        # Configure how to quantize activations.
        def get_activations_and_quantizers(self, layer):
            return [(layer.activation, tfmot.quantization.keras.quantizers.MovingAverageQuantizer(num_bits=a_bits, symmetric=False, narrow_range=False, per_axis=False))]

    # Instead of simply using quantize_annotate_model or quantize_model we must use
    # quantize_annotate_layer since it's the only one with a quantize_config argument
    def quantize_all_layers(layer):
        if isinstance(layer, tf.keras.layers.DepthwiseConv2D):
            return quantize_annotate_layer(layer, quantize_config=DepthwiseQuantizer())
        elif isinstance(layer, tf.keras.layers.Conv2D):
            return quantize_annotate_layer(layer, quantize_config=ConvQuantizer())
        return layer

    annotated_model = clone_model(
        model,
        clone_function=quantize_all_layers
    )

    with quantize_scope(
        {'Quantizer': Quantizer},
        {'ConvQuantizer': ConvQuantizer},
            {'DepthwiseQuantizer': DepthwiseQuantizer}):
        q_aware_model = quantize_apply(annotated_model)

    # *****

    # Compile and train model
    optimizer = keras.optimizers.Adam(
        learning_rate=0.001)
    q_aware_model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True),
        optimizer=optimizer, metrics=['sparse_categorical_accuracy'])

    (train_images, train_labels),_ = keras.datasets.cifar10.load_data()

    q_aware_model.fit(train_images, train_labels, batch_size=batch_size, epochs=epochs, verbose=1,
                      validation_split=0.1)

    if save:
        save_path = 'models/temp/' + name
        q_aware_model.save(save_path + '.h5')

    return q_aware_model


def temp_net():
    dropout = 0.1

    model = keras.Sequential()
    model.add(keras.layers.Conv2D(32, (3, 3), padding='same', input_shape=(32, 32, 3)))
    model.add(keras.layers.BatchNormalization())
    model.add(keras.layers.Activation('relu'))

    model.add(keras.layers.Flatten())
    model.add(keras.layers.Dense(10, activation='softmax'))

    model._name = "temp_net"

    return model


if __name__ == "__main__":
    q_model = quantization_aware_training(model=temp_net(), save=True,
                                          w_bits=8, a_bits=8, symmetric=False, narrow_range=False, per_axis=False, quantizer=QAT_ALL_VALUES, batch_size=64, epochs=1)

标签： pythontensorflowkerasquantizationquantization-aware-training