tensorflow - 使用 Estimator API 更新批量归一化均值和方差

Question

文档对此并不是 100% 清楚:

Note: when training, the moving_mean and moving_variance need to be updated. By default the update ops are placed in tf.GraphKeys.UPDATE_OPS, so they need to be added as a dependency to the train_op. For example:

(参见 https://www.tensorflow.org/api_docs/python/tf/layers/batch_normalization )

这是否意味着保存 moving_mean 和 moving_variance 所需的全部内容如下？

def model_fn(features, labels, mode, params):
   training = mode == tf.estimator.ModeKeys.TRAIN
   extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

   x = tf.reshape(features, [-1, 64, 64, 3])
   x = tf.layers.batch_normalization(x, training=training)

   # ...

  with tf.control_dependencies(extra_update_ops):
     train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())

换句话说，只需使用

with tf.control_dependencies(extra_update_ops):

要注意保存moving_mean和moving_variance吗？

Answer 1

是的，添加这些控制依赖项将保存均值和方差。