在Tensorflow/keras中的层归一化训练与评估模式

Question

我正在尝试了解训练与评估模式中的层归一化行为。我认为它类似于批处理（训练更新beta/gamma值，评估都使用其全球值）。因此，当培训= true/false时，我期望结果不同，但是结果相同。我想念什么吗？

这是工作代码：

import tensorflow as tf

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

tf.random.set_seed(42)

class MLPModel(tf.keras.Model):
    def __init__(self):
        super(MLPModel, self).__init__()
        
        self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10)
        self.norm = tf.keras.layers.LayerNormalization()
        

    def call(self, inputs):
        x = self.flatten(inputs)
        x = self.dense1(x)
        x = self.dense2(x)
        x = self.norm(x)
        return x

model = MLPModel()

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(optimizer='adam', loss=loss_fn, metrics=['accuracy'])

model.fit(x_train, y_train, epochs=2)

当我使用训练= true vs false时，结果是相同的。但是，当我用batchnorm替换层时，结果是预期的。

model(x_train[:1], training=False).numpy()

model(x_train[:1], training=True).numpy()