Keras QAT基础示例中无法将BatchNorm与Conv2D折叠问题

针对你在TensorFlow 1.15、tensorflow-model-optimization 0.6.0版本下遇到的BN层未自动折叠到Conv2D的问题，可按以下步骤处理：

• 核心原因：quantize_model仅开启量化感知训练流程，BN层折叠不会在训练过程中自动执行，必须在训练完成后手动触发权重合并操作，才能将BN的移动均值、方差等参数折叠到Conv2D的权重和偏置中。

• 具体操作步骤：

  1. 完成量化感知模型的训练（你的现有代码已完成模型定义、量化包装和编译，只需补充训练环节）。
  2. 调用tfmot.quantization.keras.fold_batch_norms函数，将训练后的模型中的BN层折叠到前置的Conv2D层。
  3. 保存折叠后的模型，此时生成的.h5文件就不会包含独立的BN层了。

修改后的完整示例代码：

import tensorflow as tf
import tensorflow_model_optimization as tfmot

smooth = 0.1  # 补充你代码中用到的label_smoothing参数值

# 原模型定义
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=(224, 224, 3)),
    tf.keras.layers.Conv2D(filters=3, kernel_size=(3, 3)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(1000)
])

# 量化包装模型
quantize_model = tfmot.quantization.keras.quantize_model
q_aware_model = quantize_model(model)

# 编译模型
q_aware_model.compile(
    optimizer='adam',
    loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=smooth),
    metrics=['accuracy']
)

# 补充训练环节（替换成你的训练数据和参数）
# q_aware_model.fit(train_dataset, epochs=10, validation_data=val_dataset)

# 训练完成后，手动折叠BN层
folded_model = tfmot.quantization.keras.fold_batch_norms(q_aware_model)

# 保存折叠后的模型
folded_model.save('quantized_folded_model.h5')

# 查看折叠后的模型结构
folded_model.summary()

• 注意事项：
  • 必须在训练完成后再执行折叠操作，因为训练过程中BN层的移动均值和方差还在动态更新，只有训练结束后这些值固定，才能正确合并到Conv2D参数中。
  • TensorFlow 1.15对应的tfmot 0.6.0版本中，fold_batch_norms仅支持Conv2D+BN的连续结构，你的Conv2D+BatchNormalization+Activation('relu')结构符合要求，会自动识别并折叠。

内容的提问来源于stack exchange，提问作者Grand