tf.distribute.get_replica_context

Возвращает текущий tf.distribute.ReplicaContext или None.

tf.distribute.get_replica_context()

Возвращает None при выполнении в межреплицируемом контексте.

Обратите внимание, что выполнение:

1. начинается в стандартном (однореплицируемом) реплицируемом контексте (эта функция вернет стандартный ReplicaContext объект);
2. переключается на межреплицируемый контекст (в этом случае эта функция вернёт None), при входе в with tf.distribute.Strategy.scope(): блок;
3. переключается на (нестандартный) реплицируемый контекст внутри strategy.experimental_run_v2(fn, ...);
4. если fn вызывает get_replica_context().merge_call(merge_fn, ...), то внутри merge_fn вы возвращаетесь в межреплицируемый контекст (и снова эта функция вернёт None).

Большинство методов tf.distribute.Strategy могут выполняться только в межреплицируемом контексте; в реплицируемом контексте следует использовать API объекта tf.distribute.ReplicaContext, возвращаемого этой функцией.

assert tf.distribute.get_replica_context() is not None  # default
with strategy.scope():
  assert tf.distribute.get_replica_context() is None

  def f():
    replica_context = tf.distribute.get_replica_context()  # for strategy
    assert replica_context is not None
    tf.print("Replica id: ", replica_context.replica_id_in_sync_group,
             " of ", replica_context.num_replicas_in_sync)

  strategy.experimental_run_v2(f)