tf.distribute.NcclAllReduce
| Просмотреть исходный код на GitHub |
Реализация NCCL all-reduce для CrossDeviceOps.
Наследуется от: CrossDeviceOps
tf.distribute.NcclAllReduce(
num_packs=1
)
Использует Nvidia NCCL для all-reduce. Для пакетного API тензоры будут переупакованы или агрегированы для более эффективной межъячейковой транспортировки.
Для операций reduce, которые не являются all-reduce, используется tf.distribute.ReductionToOneDevice.
Вот как можно использовать NcclAllReduce в tf.distribute.MirroredStrategy:
strategy = tf.distribute.MirroredStrategy( cross_device_ops=tf.distribute.NcclAllReduce())
| Аргументы | |
|---|---|
num_packs | целое неотрицательное число. Количество пакетов для разделения значений. Если ноль, упаковка не выполняется. |
| Исключения | |
|---|---|
ValueError | если num_packs отрицательно. |
Методы
batch_reduce
batch_reduce(
reduce_op, value_destination_pairs, options=None
)
Сведение значений в пункты назначения партиями.
См. tf.distribute.StrategyExtended.batch_reduce_to. Это можно вызывать только в контексте межрепликации.
| Аргументы | |
|---|---|
reduce_op | tf.distribute.ReduceOp, определяющий способ объединения значений. |
value_destination_pairs | последовательность пар (значение, пункты назначения). См. tf.distribute.CrossDeviceOps.reduce для описаний. |
options | tf.distribute.experimental.CommunicationOptions. См. tf.distribute.experimental.CommunicationOptions для подробностей. |
| Возвращаемое значение | |
|---|---|
Список tf.Tensor или tf.distribute.DistributedValues, по одному на пару в value_destination_pairs. |
| Исключения | |
|---|---|
ValueError | если value_destination_pairs не является итерируемым наборами кортежей tf.distribute.DistributedValues и пунктами назначения. |
broadcast
broadcast(
tensor, destinations
)
Рассылка tensor в destinations.
Это можно вызывать только в контексте межрепликации.
| Аргументы | |
|---|---|
tensor | объект типа tf.Tensor. Передаваемое значение. |
destinations | tf.distribute.DistributedValues, tf.Variable, tf.Tensor или строка устройства. Указывает устройства, в которые необходимо передать значение. Обратите внимание, что если это tf.Variable, значение передаётся на устройства этого объекта переменной. Этот метод не обновляет переменную. |
| Возвращаемое значение | |
|---|---|
tf.Tensor или tf.distribute.DistributedValues. |
reduce
reduce(
reduce_op, per_replica_value, destinations, options=None
)
Сведение per_replica_value в destinations.
См. tf.distribute.StrategyExtended.reduce_to. Это можно вызывать только в контексте межрепликации.
| Аргументы | |
|---|---|
reduce_op | tf.distribute.ReduceOp, определяющий способ объединения значений. |
per_replica_value | tf.distribute.DistributedValues или объект типа tf.Tensor. |
destinations | tf.distribute.DistributedValues, tf.Variable, tf.Tensor или строка устройства. Для выполнения all-reduce передайте одинаковое значение в value и destinations. Обратите внимание, что если это tf.Variable, значение сводится к устройствам этой переменной, и этот метод не обновляет переменную. |
options | tf.distribute.experimental.CommunicationOptions. См. tf.distribute.experimental.CommunicationOptions для подробностей. |
| Возвращаемое значение | |
|---|---|
tf.Tensor или tf.distribute.DistributedValues. |
| Исключения | |
|---|---|
ValueError | если per_replica_value нельзя преобразовать в tf.distribute.DistributedValues или если destinations не является строкой, tf.Variable или tf.distribute.DistributedValues. |
© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/NcclAllReduce