當前位置: 首頁>>編程示例 >>用法及示例精選 >>正文

Python tf.distribute.NcclAllReduce用法及代碼示例

NCCL all-reduce CrossDeviceOps 的實現。

繼承自：CrossDeviceOps

用法

tf.distribute.NcclAllReduce(
    num_packs=1
)

參數

  • num_packs 一個非負整數。要將值拆分成的包數。如果為零，則不進行打包。

拋出

  • ValueError 如果num_packs 為負數。

它使用 Nvidia NCCL 作為all-reduce。對於批處理 API，張量將被重新打包或聚合以更有效地 cross-device 傳輸。

對於不是 all-reduce 的歸約，它回退到 tf.distribute.ReductionToOneDevice 。

以下是在 tf.distribute.MirroredStrategy 中使用 NcclAllReduce 的方法：

strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.NcclAllReduce())

相關用法


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