TensorFlow中BatchNormalization的epsilon参数存在实现异常？

我们发现tf.keras.layers.BatchNormalization的实现与其数学模型存在不一致，问题根源可能涉及epsilon或方差参数。

问题复现步骤 #

• 初始化BN算子（命名为source_model），输入随机生成的数据集data，得到输出source_result；
• 生成扰动delta，对source_model的方差执行加delta操作，同时对epsilon执行减delta操作，得到新的BN算子follow_model；
• 将相同的data输入follow_model，得到输出follow_result；
• 计算source_result与follow_result的绝对误差和，理论上该值应趋近于0，但实际结果远大于1。

复现代码 #

# from tensorflow.keras.layers import BatchNormalization, Input
# from tensorflow.keras.models import Model, clone_model
from tensorflow._api.v1.keras.layers import BatchNormalization, Input
from tensorflow._api.v1.keras.models import Model, clone_model

import os
import re
import numpy as np

def SourceModel(shape):
    x = Input(shape=shape[1:])
    y = BatchNormalization(axis=-1)(x)
    return Model(x, y)

def FollowModel_1(source_model):
    follow_model = clone_model(source_model)
    # 读取权重
    weights = source_model.get_weights()
    weights_names = [weight.name for layer in source_model.layers for weight in layer.weights]
    variance_idx = FindWeightsIdx("variance", weights_names)

    # 扰动操作
    # delta = np.random.uniform(-1e-3, 1e-3, 1)[0]
    follow_model.layers[1].epsilon += delta     # 修改epsilon
    weights[variance_idx] -= delta
    follow_model.set_weights(weights)

    return follow_model

def FindWeightsIdx(name, weights_names):
    # 根据名称查找权重索引
    for idx, names in enumerate(weights_names):
        if re.search(name, names):
            return idx
    return -1

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

shape = (10, 32, 32, 3)
data = np.random.uniform(-1, 1, shape)
delta = -1

source_model = SourceModel(shape)
follow_model = FollowModel_1(source_model)

source_result = source_model.predict(data)
follow_result = follow_model.predict(data)
dis = np.sum(abs(source_result-follow_result))

print("delta:", delta, "; dis:", dis)

运行结果 #

delta: -1 ; dis: 4497.482

无论delta取何值，理论上误差和都应极小，但实际结果与之相悖。该问题在TensorFlow 1.x和2.x版本中均会出现。

内容的提问来源于stack exchange，提问作者Embed-Debuger