在保留稀疏性和聚类的同时从TensorFlow量化中排除重标度层

Question

我正在按照这指南对我的模型这执行量化不幸的是，我的模型包含一个无法量化的层(重标度层)。为了说明这一点，我使用quantize_annotate_layer只标记其他层进行量化。我通过调用这个代码来做到这一点：

def apply_quantization_to_non_rescaling(layer):
    if not isinstance(layer, tf.keras.layers.Rescaling):
        print('=> NOT Rescaling')
        return tfmot.quantization.keras.quantize_annotate_layer(layer, quantize_config=None)
    print('=> Rescaling')
    return layer

quant_aware_annotate_model = tf.keras.models.clone_model(
    stripped_clustered_model,
    clone_function=apply_quantization_to_non_rescaling,
)

pcqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme(preserve_sparsity=True)
)

就我的理解而言，我用quantize_annotate_layer标记所有的层，我想对其进行量化。稍后，我调用quantize_apply来执行这个量化。但是，运行此代码会导致以下错误：

=> Rescaling
=> NOT Rescaling
=> NOT Rescaling
=> NOT Rescaling
=> NOT Rescaling
=> NOT Rescaling
=> NOT Rescaling
=> NOT Rescaling
=> NOT Rescaling
=> NOT Rescaling
Traceback (most recent call last):
  File "model_2.py", line 332, in <module>
    main()
  File "model_2.py", line 304, in main
    pcqat_model = tfmot.quantization.keras.quantize_apply(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_model_optimization/python/core/keras/metrics.py", line 74, in inner
    raise error
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_model_optimization/python/core/keras/metrics.py", line 69, in inner
    results = func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py", line 474, in quantize_apply
    return keras.models.clone_model(
  File "/home/user/.local/lib/python3.8/site-packages/keras/models.py", line 453, in clone_model
    return _clone_sequential_model(
  File "/home/user/.local/lib/python3.8/site-packages/keras/models.py", line 330, in _clone_sequential_model
    if isinstance(layer, InputLayer) else layer_fn(layer))
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py", line 408, in _quantize
    full_quantize_config = quantize_registry.get_quantize_config(layer)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/cluster_preserve_quantize_registry.py", line 293, in get_quantize_config
    quantize_config = (default_8bit_quantize_registry.
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantize_registry.py", line 272, in get_quantize_config
    raise ValueError(
ValueError: `get_quantize_config()` called on an unsupported layer <class 'keras.layers.preprocessing.image_preprocessing.Rescaling'>. Check if layer is supported by calling `supports()`. Alternatively, you can use `QuantizeConfig` to specify a behavior for your layer.

输出向我显示，除第一层(即重标度层)外，所有都被标记为量化。然而，下面的错误告诉我，重标度层也用于量化。

如何从量化中排除重标度层？

更新22.4.2022:使用

pcqat_model = tfmot.quantization.keras.quantize_apply(
    quant_aware_annotate_model
)

而不是

pcqat_model = tfmot.quantization.keras.quantize_apply(
    quant_aware_annotate_model,
    tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme(preserve_sparsity=True)
)

因为这不会保持稀疏性和聚类性。

Answer 1

因此，在传递另一种量化策略时会出现错误，而不是默认的量化策略。如果您只使用pcqat_model = tfmot.quantization.keras.quantize_apply(quant_aware_annotate_model)，它就会工作。我尝试使用其他非实验量化策略，但它们也会造成一些错误。因此，如果你绝对设定了另一种策略，而不是默认的策略，这对你没有帮助，但是如果你只想使用量化，就使用默认的策略。