如何将 Tensorflow BatchNormalization 与 GradientTape 结合使用？

假设我们有一个使用 BatchNormalization 的简单 Keras 模型：

model = tf.keras.Sequential([
                     tf.keras.layers.InputLayer(input_shape=(1,)),
                     tf.keras.layers.BatchNormalization()
])

如何实际使用 GradientTape？以下似乎不起作用，因为它没有更新移动平均线？

# model training... we want the output values to be close to 150
for i in range(1000):
  x = np.random.randint(100, 110, 10).astype(np.float32)
  with tf.GradientTape() as tape:
    y = model(np.expand_dims(x, axis=1))
    loss = tf.reduce_mean(tf.square(y - 150))
  grads = tape.gradient(loss, model.variables)
  opt.apply_gradients(zip(grads, model.variables))

特别是，如果您检查移动平均值，它们将保持不变（检查 model.variables，平均值始终为 0 和 1）。我知道可以使用 .fit() 和 .predict()，但我想使用 GradientTape 并且我不知道如何执行此操作。某些版本的文档建议更新 update_ops，但这似乎在急切模式下不起作用。

特别是，经过上述训练后，以下代码将不会输出任何接近 150 的结果。

x = np.random.randint(200, 210, 100).astype(np.float32)
print(model(np.expand_dims(x, axis=1)))

使用梯度磁带模式 BatchNormalization 层应使用参数 Training=True 进行调用

example:

inp = KL.Input( (64,64,3) )
x = inp
x = KL.Conv2D(3, kernel_size=3, padding='same')(x)
x = KL.BatchNormalization()(x, training=True)
model = KM.Model(inp, x)

然后移动变量被正确更新

>>> model.layers[2].weights[2]
<tf.Variable 'batch_normalization/moving_mean:0' shape=(3,) dtype=float32, numpy
=array([-0.00062087,  0.00015137, -0.00013239], dtype=float32)>