当前位置: 首页>>代码示例 >>用法及示例精选 >>正文

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。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。