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