0

我想在 Tensorflow 2.5 中使用自定义层进行量化感知训练 (QAT)

我按照指南：

https://www.tensorflow.org/model_optimization/guide/quantization/training_comprehensive_guide

我已经创建了图层：

class RotateLayer(layers.Layer):
    def __init__(self , **kwargs):
        super(RotateLayer, self).__init__()
        filter_size = 4
        num_of_filters = 64
        self._filter_size = filter_size
        self._num_of_filters = num_of_filters

    def build(self , input_shape):
        # create filter
        n = self._filter_size ** 2 * self._num_of_filters
        w_init = tf.random_normal_initializer(stddev=np.sqrt(2. / n))
        self.kernel = []
        self.kernel.append(self.add_weight(initializer=w_init, shape=(self._filter_size,self._filter_size,1,self._num_of_filters),dtype="float32", name='kernel_rotate', trainable=True))
        self.kernel.append(self.add_weight(initializer=w_init , shape=(1, self._num_of_filters), dtype="float32",name='bias_rotate', trainable=True))
        return

    def call(self, x):
        _, H, W, _ = x.shape

        output_rotate_i = []
        for i in range(self._filter_size):
            output_rotate_j = []
            for j in range(self._filter_size):
                out = tf.nn.conv2d(x[:, i:, j:, :], tf.roll(self.kernel[0], shift=[-i, -j], axis=[1, 2]),strides=self._filter_size, padding='VALID', data_format='NHWC')
                out = tf.nn.bias_add(out , tf.squeeze(self.kernel[1]))
                output_rotate_j.append(out)
            output_rotate_i.append(tf.stack(output_rotate_j , axis=3))
        output_rotate = tf.stack(output_rotate_i , axis=2)
        output_rotate = tf.reshape(output_rotate , shape=(-1 , H , W , self._num_of_filters))
        return output_rotate

和配置：

class RotateLayerQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):
    # Configure how to quantize weights.
    def get_weights_and_quantizers(self, layer):
      return [(layer.kernel, 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, 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 {}

我已经将包含层的模型转换为量化感知：

model_to_quantize = quantize_annotate_model(model)
with quantize_scope(
  {'RotateLayerQuantizeConfig': RotateLayerQuantizeConfig,
   'RotateLayer': RotateLayer}):
    quant_aware_model = tfmot.quantization.keras.quantize_apply(model)

模型已成功转换，但是，通过检查变量和权重，我发现它们对于特定层并不相同。（对于其他未定义为自定义图层的图层，它们是相同的！）

似乎权重已重新初始化，而不是从原始模型（非量化）克隆权重。

这里我的层的变量不一样（层是0和1，0的例子）层名称：model_to_quantize.trainable_variables[0].name is 'quantize_annotate/kernel_rotate:0' quant_aware_model.trainable_variables[0].name是'rotate_layer_2/kernel_rotate:0'

    tf.abs(model_to_quantize.trainable_variables[0] - quant_aware_model.trainable_variables[0]).numpy().max()
    Out: 0.23129302

对于其他层，它是相同的：层名称：model_to_quantize.trainable_variables[2].name is 'conv2d/kernel:0' quant_aware_model.trainable_variables[2].name is 'conv2d/kernel:0'

tf.abs(model_to_quantize.trainable_variables[2] - quant_aware_model.trainable_variables[2]).numpy().max()
Out: 0.0
4

0 回答 0