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

Question

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

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))

ExponentialMovingAverage 的 TensorFlow 文档并不清楚如何在从头开始的训练循环中使用它。有人用过这个吗？

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

Answer 1

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

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

然后在优化步骤后应用 EMA。ema 对象将保留模型变量的影子变量。（您不需要调用此处，请参阅文档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")