tf.distribute.DistributedValues

Базовый класс для представления распределённых значений.

Экземпляр подкласса tf.distribute.DistributedValues создаётся при создании переменных в стратегии распределения, при итерации по tf.distribute.DistributedDataset или с помощью tf.distribute.Strategy.run. Данный базовый класс никогда не должен создаваться напрямую. tf.distribute.DistributedValues содержит значение для каждой реплики. В зависимости от подкласса, значения могут синхронизироваться при обновлении, по запросу или никогда не синхронизироваться.

Два типичных типа tf.distribute.DistributedValues — tf.types.experimental.PerReplica и tf.types.experimental.Mirrored значения.

PerReplica значения существуют на устройствах рабочих узлов, с разным значением для каждой реплики. Они создаются при итерации по распределённому набору данных, возвращаемому tf.distribute.Strategy.experimental_distribute_dataset (Пример 1, ниже) и tf.distribute.Strategy.distribute_datasets_from_function. Они также являются типичным результатом, возвращаемым tf.distribute.Strategy.run (Пример 2).

Mirrored значения похожи на PerReplica значения, за исключением того, что значения на всех репликах одинаковы. Mirrored значения синхронизируются используемой стратегией распределения, в то время как PerReplica значения не синхронизируются. Mirrored значения обычно представляют веса модели. Мы можем безопасно прочитать Mirrored значение в контексте крос-реплики, используя значение любой реплики, в то время как значения PerReplica не следует читать или изменять в контексте крос-реплики.

tf.distribute.DistributedValues может быть уменьшен с помощью strategy.reduce для получения одного значения по всем репликам (Пример 4), используемого в качестве входных данных для tf.distribute.Strategy.run (Пример 3) или для сбора значений по каждой реплике для проверки с использованием tf.distribute.Strategy.experimental_local_results (Пример 5).

Примеры использования:

Создан из tf.distribute.DistributedDataset:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
distributed_values = next(dataset_iterator)
distributed_values
PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=float32, numpy=array([5.], dtype=float32)>,
  1: <tf.Tensor: shape=(1,), dtype=float32, numpy=array([6.], dtype=float32)>
}

Возвращается run:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
@tf.function
def run():
  ctx = tf.distribute.get_replica_context()
  return ctx.replica_id_in_sync_group
distributed_values = strategy.run(run)
distributed_values
PerReplica:{
  0: <tf.Tensor: shape=(), dtype=int32, numpy=0>,
  1: <tf.Tensor: shape=(), dtype=int32, numpy=1>
}

В качестве входных данных для run:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
distributed_values = next(dataset_iterator)
@tf.function
def run(input):
  return input + 1.0
updated_value = strategy.run(run, args=(distributed_values,))
updated_value
PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=float32, numpy=array([6.], dtype=float32)>,
  1: <tf.Tensor: shape=(1,), dtype=float32, numpy=array([7.], dtype=float32)>
}

В качестве входных данных для reduce:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
distributed_values = next(dataset_iterator)
reduced_value = strategy.reduce(tf.distribute.ReduceOp.SUM,
                                distributed_values,
                                axis = 0)
reduced_value
<tf.Tensor: shape=(), dtype=float32, numpy=11.0>

Как проверить значения по каждой реплике локально:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset))
per_replica_values = strategy.experimental_local_results(
   distributed_values)
per_replica_values
(<tf.Tensor: shape=(1,), dtype=float32, numpy=array([5.], dtype=float32)>,
 <tf.Tensor: shape=(1,), dtype=float32, numpy=array([6.], dtype=float32)>)

© 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/api_docs/python/tf/distribute/DistributedValues