TensorFlow模型优化中，如何正确保存量化封装模型的检查点？

问题根源 #

你遇到的问题是因为quantize_model包装后的模型包含量化感知训练（QAT）相关的非可训练变量（比如量化层的移动均值、方差，以及量化阈值参数），而默认的save_weights("qat_weights.h5")只会保存模型的可训练权重，这些QAT关键变量没有被保存；即使后续加载权重，QAT层的状态也会被重置，导致训练看起来像是从头开始。

可行解决方案 #

方案1：保存/加载完整模型（推荐） #

直接保存整个QAT模型（包含所有结构、权重和QAT状态），是最稳妥的方式：

保存模型

from tensorflow_model_optimization.quantization.keras import quantize_model

model = define_model()
qat_model = quantize_model(model)
qat_model.fit(...)
# 保存为SavedModel格式（比H5更适配TF2.x）
qat_model.save("qat_trained_model")

加载模型

加载时需要指定QAT相关的自定义对象，否则模型无法正确解析量化层：

import tensorflow as tf
import tensorflow_model_optimization as tfmot

# 直接加载完整模型，包含所有QAT状态
qat_model = tf.keras.models.load_model(
    "qat_trained_model",
    custom_objects={'QuantizeWrapper': tfmot.quantization.keras.QuantizeWrapper}
)
# 此时调用fit会延续之前的训练进度
qat_model.fit(...)

方案2：手动保存/加载所有变量（若必须用权重文件） #

如果坚持使用save_weights，需要确保保存所有变量（包括非可训练的），并且加载前先初始化模型的所有变量：

保存权重

改用TensorFlow格式保存，支持保存所有变量：

qat_model.fit(...)
# 保存为TF格式权重（会保存所有变量，包括非可训练的QAT参数）
qat_model.save_weights("qat_full_weights", save_format='tf')

加载权重

加载前必须先构建模型，确保所有QAT变量被初始化：

model = define_model()
qat_model = quantize_model(model)
# 先构建模型，替换为你的模型实际输入形状
qat_model.build(input_shape=(None, 224, 224, 3))
# 加载所有权重，包括QAT相关的非可训练变量
qat_model.load_weights("qat_full_weights")
# 此时训练会延续之前的进度
qat_model.fit(...)

额外检查点 #

如果权重加载正确但训练进度显示从0%开始，确认fit时是否设置了initial_epoch参数：比如之前训练了50个epoch，加载后需要设置qat_model.fit(..., initial_epoch=50)，这样进度条会从对应位置开始显示。

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