tf.distribute.experimental.CollectiveHints

Подсказки для коллективных операций, таких как AllReduce.

tf.distribute.experimental.CollectiveHints(
    bytes_per_pack=0, timeout_seconds=None
)

Это можно передать в методы, такие как tf.distribute.get_replica_context().all_reduce(), чтобы оптимизировать производительность коллективной операции. Обратите внимание, что это только подсказки, которые могут или не могут изменить фактическое поведение. Некоторые параметры применяются только к определённой стратегии и игнорируются другими.

Одной из распространённых оптимизаций является разделение all-reduce градиентов на несколько пакетов, чтобы обновления весов могли перекрываться с all-reduce градиентов.

Примеры:

• bytes_per_pack

hints = tf.distribute.experimental.CollectiveHints(
    bytes_per_pack=50 * 1024 * 1024)
grads = tf.distribute.get_replica_context().all_reduce(
    'sum', grads, experimental_hints=hints)
optimizer.apply_gradients(zip(grads, vars),
    experimental_aggregate_gradients=False)

• timeout_seconds

strategy = tf.distribute.MirroredStrategy()
hints = tf.distribute.experimental.CollectiveHints(
    timeout_seconds=120.0)
try:
  strategy.reduce("sum", v, axis=None, experimental_hints=hints)
except tf.errors.DeadlineExceededError:
  do_something()