0

我正在关注量化意识培训综合指南，并在 QAT 中为自定义层而苦苦挣扎，tf=2.6.0使用py=3.9.7. 下面是我的问题的一个玩具示例：

我写了一个简单的自定义层来实现 Conv2D

class MyConv(tf.keras.layers.Layer):
'''costume conv2d'''

def __init__(self, filt=1, name=None, **kwargs):
    super(MyConv, self).__init__(name=name)
    self.filt = filt
    super(MyConv, self).__init__(**kwargs)

def get_config(self):
    config = super().get_config().copy()
    config.update({"filt": self.filt})
    return config

def build(self, shape):
    self.conv = tf.keras.layers.Conv2D(self.filt, 1, padding="same")

def call(self, input):
    return self.conv(input)

我用那个层创建了一个小模型，然后递归地传递它的层并使用注释tfmot.guantization.keras.quantize_annotate_layer（每个自定义层可能有更多需要量化的自定义子层）。然后我适用tfmot.quantization.keras.quantize_apply于带注释的模型。结果模型由所有量化层组成，除了我的自定义层之外，还没有被量化。

附模型摘要

我会注意到，当我用MyConv下面的代码替换自定义层时，如综合指南中所示，量化是有效的。

def MyConv(tf.keras.layers.Conv2D):
    pass

请帮我解决这个问题。可能是我的问题QuantizeConfig？

以下是我的完整代码：

import tensorflow as tf
import tensorflow_model_optimization as tfmot

class MyConv(tf.keras.layers.Layer):
'''costume conv2d'''

def __init__(self, filt=1, name=None, **kwargs):
    super(MyConv, self).__init__(name=name)
    self.filt = filt
    super(MyConv, self).__init__(**kwargs)

def get_config(self):
    config = super().get_config().copy()
    config.update({"filt": self.filt})
    return config

def build(self, shape):
    self.conv = tfmot.quantization.keras.quantize_annotate_layer(tf.keras.layers.Conv2D(self.filt, 1, padding="same"))

def call(self, input):
    return self.conv(input)


def get_toy_model():
  input = tf.keras.Input((10, 10, 1), name='input')
  x = tf.keras.layers.Conv2D(1, 3, padding="same")(input)
  x = tf.keras.layers.ReLU()(x)
  x = MyConv()(x)
  for _ in range(2):
      y = tf.keras.layers.Conv2D(1, 3, padding="same")(x)
      y = tf.keras.layers.ReLU()(y)
  out = tf.keras.layers.Conv2D(1, 3, padding="same")(y)
  return tf.keras.Model(input, out, name='toy_Conv2D')

LastValueQuantizer = tfmot.quantization.keras.quantizers.LastValueQuantizer
MovingAverageQuantizer = tfmot.quantization.keras.quantizers.MovingAverageQuantizer


class DefaultCostumeQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):
    # Configure how to quantize weights.
    def get_weights_and_quantizers(self, layer):
        return []

    # Configure how to quantize activations.
    def get_activations_and_quantizers(self, layer):
        return []

    def set_quantize_weights(self, layer, quantize_weights):
        pass

    def set_quantize_activations(self, layer, quantize_activations):
        pass

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

    def get_config(self):
        return {}


def recursive_depth_layers(layer):
  for l in list(layer.__dict__.values()):
    if isinstance(l, tf.keras.layers.Layer):
        recursive_depth_layers(l)
        if isinstance(l, (
        tf.keras.layers.Dense, tf.keras.layers.Conv2D, tf.keras.layers.ReLU, tf.keras.layers.LeakyReLU, tf.keras.layers.Activation)):
            ql = tfmot.quantization.keras.quantize_annotate_layer(l, DefaultCostumeQuantizeConfig())
            ql._name += "_" + l.name
            return ql


def apply_quantization(layer):
    # regular layer
    if isinstance(layer, (tf.keras.layers.Dense, tf.keras.layers.Conv2D, tf.keras.layers.ReLU, tf.keras.layers.LeakyReLU,tf.keras.layers.Activation)):
      l = tfmot.quantization.keras.quantize_annotate_layer(layer, DefaultCostumeQuantizeConfig())
      l._name += '_' + layer.name
      return l
    if layer.__module__ == "__main__":
    # custom layer
      recursive_depth_layers(layer)
      l = tfmot.quantization.keras.quantize_annotate_layer(layer, DefaultCostumeQuantizeConfig())
      l._name += '_' + layer.name
      return l
  return layer

model = get_toy_model()
model.summary()

annotated_model = tf.keras.models.clone_model(model, clone_function=apply_quantization)
annotated_model.summary()

quantize_scope = tfmot.quantization.keras.quantize_scope
with quantize_scope({'DefaultCostumeQuantizeConfig': DefaultCostumeQuantizeConfig, 'MyConv': MyConv}):  
    quant_aware_model = tfmot.quantization.keras.quantize_apply(annotated_model)
    quant_aware_model._name += "_quant"
quant_aware_model.summary()
quant_aware_model.compile()
4

0 回答 0