Как получить квантованные веса из обучения квантованию TensorFlow с экспериментальным квантованием

Я использую обучающий API с поддержкой квантования TensorFlow и хочу развернуть модель с произвольной разрядностью. Поскольку для развертывания tflite поддерживается только 8-битное квантование, я буду развертывать его с помощью специального алгоритма вывода, но мне все еще нужно получить доступ к весам модели в правильном размере.

В настоящее время после использования обучения с учетом квантования моя модель все еще находится в формате с плавающей запятой, и, насколько я видел, единственный способ получить доступ к квантованным весам - это преобразовать модель в формат tflite. Однако это невозможно при использовании экспериментальных функций.

Вот мой класс конфигурации квантования:

    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 ModifiedQuantizer(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=bits, symmetric=symmetric, narrow_range=narrow_range, per_axis=per_axis))]

А вот как я квантую модель:

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

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 ModifiedQuantizer(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=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=bits, symmetric=False, narrow_range=False, per_axis=False))]

    def quantize_all_layers(layer):
        for supported_layer in supported_layers:
            if isinstance(layer, supported_layer):
                return quantize_annotate_layer(layer, quantize_config=ModifiedQuantizer())
        # print(layer.name)
        return layer
    annotated_model = clone_model(
        model,
        clone_function=quantize_all_layers
    )

with quantize_scope(
    {'Quantizer': Quantizer},
    {'ModifiedQuantizer': ModifiedQuantizer},
    {'_relu6': models._relu6}):
    q_aware_model = quantize_apply(annotated_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, val_images, val_labels, _, _ = cifar10.load()

q_aware_model.fit(train_images, train_labels, batch_size=64, epochs=1, verbose=1,
                  validation_data=(val_images, val_labels))

Ранее было сказано, при использовании, например, bits = 4 в ModifiedQuantizer, модель по-прежнему сохраняется с плавающей запятой, и я не знаю, как получить доступ к квантованным весам.

Спасибо!


python keras tensorflow quantization quantization-aware-training
person LucasStromberg    schedule 03.03.2021    source источник

Ответы (1)


arrow_upward
0
arrow_downward

Я подозреваю, что вы можете получить квантованные веса, вызвав LastValueQuantizer.__call__ на тензоре весов данного слоя. Вопрос в том, как вызвать этот метод.

Текущая подпись:

    LastValueQuantizer.__call__(inputs, training, weights, **kwargs)

Я предполагаю, что inputs - это веса слоя, а weights - это значение, возвращаемое LastValueQuantizer.build. Если бы вы могли получить ссылку на weights, возвращаемый build, я бы надеялся, что было бы просто квантовать веса слоя напрямую, используя LastValueQuantizer.__call__.

[nav] In [1]: from tensorflow_model_optimization.quantization.keras.quantizers import LastValueQuantizer
INFO:tensorflow:Enabling eager execution
INFO:tensorflow:Enabling v2 tensorshape
INFO:tensorflow:Enabling resource variables
INFO:tensorflow:Enabling tensor equality
INFO:tensorflow:Enabling control flow v2

[nav] In [2]: q = LastValueQuantizer(num_bits=3, per_axis=True, symmetric=True, narrow_range=True)

[ins] In [3]: ??q.__call__
Signature: q.__call__(inputs, training, weights, **kwargs)
Source:   
  def __call__(self, inputs, training, weights, **kwargs):
    """Quantize tensor.

    Args:
      inputs: Input tensor to be quantized.
      training: Whether the graph is currently training.
      weights: Dictionary of weights the quantizer can use to quantize the
        tensor. This contains the weights created in the `build` function.
      **kwargs: Additional variables which may be passed to the quantizer.

    Returns:
      Quantized tensor.
    """
    return quant_ops.LastValueQuantize(
        inputs,
        weights['min_var'],
        weights['max_var'],
        is_training=training,
        num_bits=self.num_bits,
        per_channel=self.per_axis,
        symmetric=self.symmetric,
        narrow_range=self.narrow_range
    )
person ndronen    schedule 29.03.2021
comment
Спасибо. Я займусь этим. В настоящее время я квантовал каждый слой, вычислив масштаб и нулевую точку для каждого тензора веса. Однако, насколько мне известно, TensorFlow никогда не указывает точную реализацию получения параметров квантования, так что это довольно временное решение. - person LucasStromberg; 31.03.2021