TensorFlow中BatchNormalization使用及模型恢复问题求助

我之前在TensorFlow里用Batch Norm恢复参数时也踩过一模一样的坑，给你几个实用的排查和解决思路：

1. 先确认变量名称完全匹配 #

手动恢复moving_mean和moving_variance时，最容易犯的错就是变量名称不匹配。比如你用tf.layers.batch_normalization定义的层，变量名通常是[层名]/moving_mean:0、[层名]/moving_variance:0这种格式。你可以先打印出checkpoint里的所有变量，对比模型中当前变量的名称：

from tensorflow.python.tools.inspect_checkpoint import print_tensors_in_checkpoint_file
print_tensors_in_checkpoint_file(file_name='你的模型路径.ckpt', tensor_name='', all_tensors=True)

一定要确保你加载的变量名和模型里的变量名完全一致，哪怕多一个前缀或者少个下划线都会导致恢复失败。

2. 别混淆训练/推理模式 #

Batch Norm在训练和推理时的逻辑完全不同：训练时会用当前batch的均值方差，同时更新滑动统计量；推理时才会用保存的moving_mean和moving_variance。如果你恢复参数后，不小心在推理时还设置了training=True，那模型会立刻更新滑动统计量，把你刚恢复的值覆盖掉！

所以恢复后做推理时，一定要给Batch Norm层传入training=False：

# 假设你的模型定义是这样的
def build_model(inputs, training):
    x = tf.layers.batch_normalization(inputs, training=training)
    # ...其他层定义

# 恢复参数后的推理代码
with tf.Session() as sess:
    # 执行你的恢复操作
    sess.run(assign_values_to_batchNorm())
    # 推理时必须设training=False
    results = sess.run(output_tensor, feed_dict={
        input_placeholder: 你的数据,
        training_placeholder: False
    })

3. 确保assign操作正确执行 #

你提到调用了assign_values_to_batchNorm()，要确认这个函数里的assign操作是正确绑定到对应的变量上的，而且在会话中确实执行了。可以把所有assign操作打包成列表，一次性run，之后还可以验证一下恢复结果：

def assign_values_to_batchNorm():
    # 假设你已经从checkpoint加载了mean_val和var_val
    assign_mean_op = tf.assign(模型中的moving_mean变量, mean_val)
    assign_var_op = tf.assign(模型中的moving_variance变量, var_val)
    return [assign_mean_op, assign_var_op]

with tf.Session() as sess:
    # 先初始化所有全局变量，再用assign覆盖Batch Norm的统计量
    sess.run(tf.global_variables_initializer())
    # 执行恢复操作
    sess.run(assign_values_to_batchNorm())
    # 验证是否恢复成功
    restored_mean, restored_var = sess.run([模型中的moving_mean变量, 模型中的moving_variance变量])
    print("恢复后的moving_mean:", restored_mean)
    print("恢复后的moving_variance:", restored_var)

先初始化再覆盖，能避免变量未初始化的问题，也能确认恢复是否真的生效。

4. 检查是否用了滑动平均的旧接口 #

如果你用的是tf.contrib.layers.batch_norm或者配合了tf.train.ExponentialMovingAverage，那恢复逻辑会不一样。比如用ExponentialMovingAverage时，滑动统计量会被存在带有ExponentialMovingAverage后缀的变量里，这时候不需要手动assign，直接用Saver加载滑动平均后的变量就行：

ema = tf.train.ExponentialMovingAverage(decay=0.99)
# 获取所有需要恢复的滑动平均变量
var_list = ema.variables_to_restore()
saver = tf.train.Saver(var_list)
saver.restore(sess, '你的模型路径.ckpt')

这种情况下Saver会自动把moving_mean和moving_variance的滑动平均值恢复到对应的变量里。

内容的提问来源于stack exchange，提问作者I. A