Результаты нормализации партии несовместимы между Keras и Tensorflow для примера игрушек: почему?

Question

Я пытаюсь заставить BatchNormalization правильно работать в Keras (2.2.4), но мне не повезло.Его поведение кажется непоследовательным в model.fit() и model.predict() / evaluate().

Моя первоначальная проблема была в контексте сложной установки GAN с различными слоями, которые переключались между замороженными и незамерзшими.Пытаясь понять, почему что-то не так, я создал этот игрушечный пример, в котором есть только один слой BatchNormalization, пытающийся изучить функцию идентификации, и больше нет бессмысленных / замерзающих бессмыслиц:

import numpy as np
import keras
from keras.layers import Input, BatchNormalization

if __name__ == '__main__':
    x = np.random.uniform(low=-1.0, high=1.0, size=(20,1,))
    y = x

    ip = tx =Input(shape=[1,])
    tx = BatchNormalization()(tx)
    mod = keras.Model(inputs=ip, outputs=tx)

    mod.compile(optimizer=keras.optimizers.SGD(lr=0.01), loss="mse")
    mod.fit(x, y, epochs=2000)
    print("mod evaluate", mod.evaluate(x, y))

Я также создал реализацию чистого тензорного потока, которая пытается сделать то же самое:

import tensorflow as tf
import numpy as np

if __name__ == '__main__':
    x = np.random.uniform(low=-1.0, high=1.0, size=(20,1))
    y = x

    input = tx = tf.Variable(initial_value=x)
    is_training = tf.Variable(initial_value=False)
    tx = tf.layers.batch_normalization(tx, training=is_training)
    output = tx
    train_output = tf.placeholder(tf.float32, shape=(None, 1))

    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    loss = tf.losses.mean_squared_error(train_output, output)
    train_op = tf.train.MomentumOptimizer(learning_rate=0.01, momentum=0.9).minimize(loss)
    train_op = tf.group([train_op, update_ops])

    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    is_training.load(False, sess)
    py = sess.run(output)
    print("py", py)
    print("mse", np.mean(np.square(py - x)))

    print("training...")
    for i in range(2000):
        is_training.load(True, sess)
        sess.run(train_op, {train_output: y})

        is_training.load(False, sess)
        py = sess.run(output)
        print("step", i, "mse", np.mean(np.square(py - x)))

    is_training.load(False, sess)
    py = sess.run(output)
    print("mse", np.mean(np.square(py - x)))

Я ожидаю, что результаты Keras и Tensorflow будут похожими.Но код Keras выводит ошибку MSE в конце примерно 0,0002, показывая при этом потерю примерно 1e-12 во время тренировки.Код Tensorflow дает гораздо меньшую ошибку, примерно 1e-16 для обучения и тестирования.Результат Tensorflow является ожидаемым, так как изучение функции тождества должно быть тривиальным, и как только скользящее среднее и дисперсия BN сходятся, результаты должны быть идентичными для всего поезда / теста.

Почему это происходит?Почему поведение Tensorflow и Keras несовместимо, и как можно получить меньшую ошибку, используя чистые Keras?Любое понимание ситуации приветствуется.