如何在自定义 TF2.4 训练循环中使用指数移动平均线

我有一个自定义训练循环，可以简化如下

inputs = tf.keras.Input(dtype=tf.float32, shape=(None, None, 3))
model = tf.keras.Model({"inputs": inputs}, {"loss": f(inputs)})
optimizer = tf.keras.optimizers.SGD(learning_rate=0.1, momentum=0.9, nesterov=True)

for inputs in batches:
    with tf.GradientTape() as tape:
        results = model(inputs, training=True)
    grads = tape.gradient(results["loss"], model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))

The ExponentialMovingAverage 的 TensorFlow 文档 https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage不清楚应该如何使用它从头开始训练循环 https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch。有人用过这个吗？

此外，如果影子变量仍然在内存中，应该如何将影子变量恢复到模型中，以及如何检查训练变量是否已正确更新？

在训练循环之前创建 EMA 对象：

ema = tf.train.ExponentialMovingAverage(decay=0.9999)

然后在优化步骤后应用 EMA。 ema 对象将保留模型变量的影子变量。 （您不需要致电tf.control_dependencies在这里，请参阅注释中的文档 https://www.tensorflow.org/api_docs/python/tf/control_dependencies)

optimizer.apply_gradients(zip(grads, model.trainable_variables))
ema.apply(model.trainable_variables)

然后，在模型中使用影子变量的一种方法可能是通过调用将影子变量分配给模型的变量averageEMA 对象的方法：

for var in model.trainable_variables:
    var.assign(ema.average(var))
model.save("model_with_shadow_variables.h5")