如何在量化感知训练后合并 ReLU [英] How to merge ReLU after quantization aware training

问题描述

我有一个包含 Conv2D 层和 ReLU 激活的网络，声明如下:

I have a network which contains Conv2D layers followed by ReLU activations, declared as such:

x = layers.Conv2D(self.hparams['channels_count'], kernel_size=(4,1))(x)
x = layers.ReLU()(x)

并移植到 TFLite 中，表示如下:

And it is ported to TFLite with the following representaiton:

没有 Q-aware 训练的基本 TFLite 网络

但是，在对网络执行量化感知训练并再次移植之后，ReLU 层现在在图中是明确的:

However, after performing quantization-aware training on the network and porting it again, the ReLU layers are now explicit in the graph:

Q-aware 训练后的 TFLite 网络

这导致它们在目标上被单独处理，而不是在 Conv2D 内核评估期间，导致我的整个网络性能损失 10%.

This results in them being processed separately on the target instead of during the evaluation of the Conv2D kernel, inducing a 10% performance loss in my overall network.

使用以下隐式语法声明激活不会产生问题:

Declaring the activation with the following implicit syntax does not produce the problem:

x = layers.Conv2D(self.hparams['channels_count'], kernel_size=(4,1), activation='relu')(x)

具有隐式 ReLU 激活的基本 TFLite 网络

Q-aware 训练后带有隐式 ReLU 的 TFLite 网络

然而，这将网络限制为基本的 ReLU 激活，而我想使用不能以这种方式声明的 ReLU6.

However, this restricts the network to basic ReLU activation, whereas I would like to use ReLU6 which cannot be declared in this way.

这是 TFLite 问题吗?如果没有，有没有办法防止ReLU层被分裂?或者，有没有办法在量化感知训练后手动将 ReLU 层合并回 Conv2D 层?

Is this a TFLite issue? If not, is there a way to prevent the ReLU layer from being split? Or alternatively, is there a way to manually merge the ReLU layers back into the Conv2D layers after the quantization-aware training?

质量检查培训代码:

def learn_qaware(self):
quantize_model = tfmot.quantization.keras.quantize_model
self.model = quantize_model(self.model)

training_generator = SCDataGenerator(self.training_set)
validate_generator = SCDataGenerator(self.validate_set)

self.model.compile(
            optimizer=self.configure_optimizers(qa_learn=True),
            loss=self.get_LLP_loss(),
            metrics=self.get_metrics(),
            run_eagerly=config['eager_mode'],
        )
self.model.fit(
    training_generator,
    epochs = self.hparams['max_epochs'],
    batch_size = 1,
    shuffle = self.hparams['shuffle_curves'],
    validation_data = validate_generator,
    callbacks = self.get_callbacks(qa_learn=True),
)

量化的TFLite模型生成代码:

Quantized TFLite model generation code:

def tflite_convert(classifier):
output_file = get_tflite_filename(classifier.model_path)

# Convert the model to the TensorFlow Lite format without quantization
saved_shape = classifier.model.input.shape.as_list()
fixed_shape = saved_shape
fixed_shape[0] = 1
classifier.model.input.set_shape(fixed_shape) # Force batch size to 1 for generation
converter = tf.lite.TFLiteConverter.from_keras_model(classifier.model)
classifier.model.input.set_shape(saved_shape)

# Set the optimization flag.
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# Enforce integer only quantization
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
# Provide a representative dataset to ensure we quantize correctly.
if config['eager_mode']:
    tf.executing_eagerly()
def representative_dataset():
    for x in classifier.validate_set.get_all_inputs():
        rs = x.reshape(1, x.shape[0], 1, 1).astype(np.float32)
        yield([rs])

converter.representative_dataset = representative_dataset
model_tflite = converter.convert()

# Save the model to disk
open(output_file, "wb").write(model_tflite)

return TFLite_model(output_file)

推荐答案

我找到了一种解决方法，该方法通过实例化模型的非训练版本，然后在转换为 TFLite 之前复制量化感知训练模型的权重.

I have found a workaround which works by instantiating a non-trained version of the model, then copying over the weights from the quantization aware trained model before converting to TFLite.

这似乎是一个黑客，所以我仍在寻找更清洁的解决方案.

This seems like quite a hack, so I'm still on the lookout for a cleaner solution.

解决方法的代码:

def dequantize(self):
    if not hasattr(self, 'fp_model') or not self.fp_model:
        self.fp_model = self.get_default_model()

    def find_layer_in_model(name, model):
        for layer in model.layers:
            if layer.name == name:
                return layer
        return None

    def find_weight_group_in_layer(name, layer):
        for weight_group in quant_layer.trainable_weights:
            if weight_group.name == name:
                return weight_group
        return None

    for layer in self.fp_model.layers:
        if 'input' in layer.name or 'quantize_layer' in layer.name:
            continue
        
        QUANT_TAG = "quant_"
        quant_layer = find_layer_in_model(QUANT_TAG+layer.name,self.model)
        if quant_layer is None:
            raise RuntimeError('Failed to match layer ' + layer.name)

        for i, weight_group in enumerate(layer.trainable_weights):
            quant_weight_group = find_weight_group_in_layer(QUANT_TAG+weight_group.name, quant_layer)
            if quant_weight_group is None:
                quant_weight_group = find_weight_group_in_layer(weight_group.name, quant_layer)
                if quant_weight_group is None:
                    raise RuntimeError('Failed to match weight group ' + weight_group.name)

            layer.trainable_weights[i].assign(quant_weight_group)

    self.model = self.fp_model

这篇关于如何在量化感知训练后合并 ReLU的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！