如何将QAT量化感知训练的TensorFlow模型转为TFLite？异常排查与解决

我正在使用QAT（量化感知训练）量化神经网络，并希望将其转换为TFLite模型。量化节点已添加到模型骨架图中生成新计算图，我能够在新图上加载训练好的QAT权重，也能生成TFLite模型，但输出结果异常：图像变暗，仿佛神经网络抑制了输入，难以定位错误来源。

我的代码如下：

model = skeleton_model() # 创建模型骨架

def apply_quantization(layer): 
    if isinstance(layer, tf.keras.layers.Conv2DTranspose):
        return tfmot.quantization.keras.quantize_annotate_layer(layer, quantize_config=CustomQuantizeConfig())
    elif isinstance(layer, tf.keras.layers.Conv2D):
        # 用默认量化配置标注Conv2D层
        return tfmot.quantization.keras.quantize_annotate_layer(layer)
    
    return layer

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

# with quantize_scope({'CustomQuantizeConfig': CustomQuantizeConfig}):

qat_model = tfmot.quantization.keras.quantize_apply(annotated_model)
# 标注后的模型包含加载QAT模型所需的所有量化节点（模型经过QAT后基础骨架结构会改变）

# 加载QAT checkpoint
checkpoint_path = os.path.join(ckpt_dir, "cp-40.ckpt")
qat_model.load_weights(checkpoint_path)
converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)

# converter.optimizations = [tf.lite.Optimize.DEFAULT]
quantized_tflite_model = converter.convert()

问题解决方法 #

代码修改 #

仅需一处小修改即可正确加载QAT权重并完成TFLite转换：此前调用quantize_apply时未放入quantize_scope，导致量化效果错误，将其放入作用域即可得到正确结果：

with quantize_scope({'CustomQuantizeConfig': CustomQuantizeConfig}):
    qat_model = tfmot.quantization.keras.quantize_apply(annotated_model)

版本兼容性注意事项 #

• 转换对TensorFlow版本敏感：使用自定义量化配置量化Conv2DTranspose层时，最初用tf2.3搭配TFLite2.10无法完成正确转换，改用tf2.10/tf2.14搭配TFLite2.18则可成功转换。
• tf>=2.15会出现Keras兼容错误：新版本Keras无法与旧版Keras节点兼容，需使用新版Keras节点训练的模型。

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