Python tf.distribute.experimental.CommunicationOptions用法及代碼示例

跨設備通信的選項，例如All-reduce。

用法

tf.distribute.experimental.CommunicationOptions(
    bytes_per_pack=0, timeout_seconds=None,
    implementation=tf.distribute.experimental.CollectiveCommunication.AUTO
)

參數

bytes_per_pack 一個非負整數。將集體操作分成一定大小的包。如果為零，則自動確定該值。除了 TPUStrategy 之外，所有 multi-replica 策略都遵循此提示。
timeout_seconds 浮點數或無，以秒為單位超時。如果不是 None，如果花費的時間超過此超時時間，則集體會引發 tf.errors.DeadlineExceededError。零禁用超時。這在調試掛起問題時很有用。這應該隻用於調試，因為它為每個集合創建一個新線程，即 timeout_seconds * num_collectives_per_second 更多線程的開銷。這僅適用於 tf.distribute.experimental.MultiWorkerMirroredStrategy 。
implementation 一個 tf.distribute.experimental.CommunicationImplementation 。這是對首選通信實現的提示。可能的值包括 AUTO , RING 和 NCCL 。 NCCL 通常對 GPU 性能更高，但不適用於 CPU。這僅適用於 tf.distribute.experimental.MultiWorkerMirroredStrategy 。

拋出

ValueError 當參數具有無效值時。

這可以傳遞給 tf.distribute.get_replica_context().all_reduce() 等方法以優化集體操作性能。請注意，這些隻是提示，可能會也可能不會改變實際行為。某些選項僅適用於某些策略，而被其他選項忽略。

一種常見的優化是將梯度all-reduce分成多個包，以便權重更新可以與梯度all-reduce重疊。

例子：

options = tf.distribute.experimental.CommunicationOptions(
    bytes_per_pack=50 * 1024 * 1024,
    timeout_seconds=120.0,
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
)
grads = tf.distribute.get_replica_context().all_reduce(
    'sum', grads, options=options)
optimizer.apply_gradients(zip(grads, vars),
    experimental_aggregate_gradients=False)

相關用法

注：本文由純淨天空篩選整理自tensorflow.org大神的英文原創作品 tf.distribute.experimental.CommunicationOptions。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。