如何在tf.keras中调用BatchNormalization时传入已有的统计参数

首先明确结论：你查到的“必须设置training=True完成训练才能在推理模式下使用BatchNormalization层”的说法是错误的。tf.keras原生支持手动传入预训练的统计参数和缩放平移参数，无需额外训练即可直接用于推理场景，具体操作方法如下： #

操作步骤 #

• 第一步：初始化BN层时匹配你的数据格式配置参数
  需重点确认axis参数和输入的通道维度匹配：输入为通道最后格式(NHWC)时axis=-1，输入为通道优先格式(NCHW)时axis=1；如果你的预训练参数包含beta和gamma，保持默认的center=True、scale=True即可，不需要这两个参数的话对应设为False。训练相关的momentum等参数无需配置，推理场景不会用到。
• 第二步：手动给层赋值预训练参数
  tf.keras的BatchNormalization层权重顺序固定为[gamma, beta, moving_mean, moving_variance]，调用set_weights方法传入参数列表即可，注意所有参数的shape要和通道数一致。
• 第三步：推理调用时指定training=False
  该模式下BN层会直接使用你传入的移动均值、移动方差做归一化，再用gamma和beta做缩放平移，不会更新任何内部统计量。

代码示例 #

import tensorflow as tf

# 替换为你的预训练参数，示例通道数为64
pretrained_gamma = tf.ones(shape=(64,))
pretrained_beta = tf.zeros(shape=(64,))
pretrained_moving_mean = tf.random.normal(shape=(64,))
pretrained_moving_var = tf.random.uniform(shape=(64,), minval=0.5, maxval=1.5)

# 初始化BN层，示例为通道最后格式
bn_layer = tf.keras.layers.BatchNormalization(axis=-1, center=True, scale=True)
# 先build层指定输入shape，对应格式为(批次, 高, 宽, 通道数)
bn_layer.build(input_shape=(None, None, None, 64))

# 按顺序传入权重
bn_layer.set_weights([pretrained_gamma, pretrained_beta, pretrained_moving_mean, pretrained_moving_var])

# 推理调用
test_input = tf.random.normal(shape=(2, 224, 224, 64))
test_output = bn_layer(test_input, training=False)

注意事项 #

如果你在初始化时关闭了center或scale，权重列表长度会对应减少，比如center=False, scale=True时权重顺序为[gamma, moving_mean, moving_variance]，需要对应调整传入set_weights的参数顺序，否则计算结果会完全错误。

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