量化感知训练后如何合并ReLU

Question

我有一个网络，其中包含Conv2D层，然后是ReLU激活，声明如下：

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

并将其移植到具有以下代表性的TFLite上：

没有q感知训练的基本TFLite网络

但是，在网络上执行量化感知培训并再次移植之后，ReLU层现在在图中显式显示：

Q感知训练后的TFLite网络

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

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

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

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

Q感知训练后的内隐TFLite网络

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

这是TFLite问题吗？如果没有，是否有办法防止ReLU层被分割？或者，在量化感知培训之后，是否有一种手动将ReLU层合并回Conv2D层的方法？

编辑： QA培训代码：

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模型生成代码：

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)

Answer 1

我找到了一个解决办法，通过实例化一个未经训练的模型，然后在转换为TFLite之前，从量化感知训练模型中复制权重。

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

解决方案的代码：

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