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