tf.distribute.OneDeviceStrategy

Стратегия распределения для выполнения на одном устройстве.

Наследуется от: Strategy

tf.distribute.OneDeviceStrategy(
    device
)

Используется в блокнотах

При использовании этой стратегии все переменные, созданные в ее области видимости, будут размещены на указанном устройстве. Входные данные, распределенные с помощью этой стратегии, будут предварительно загружены на указанное устройство. Кроме того, все функции, вызываемые через strategy.run, также будут размещены на указанном устройстве.

Эту стратегию можно использовать для тестирования кода с API tf.distribute.Strategy перед переходом к другим стратегиям, которые действительно распределяют вычисления между несколькими устройствами или машинами.

Например:

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

with strategy.scope():
  v = tf.Variable(1.0)
  print(v.device)  # /job:localhost/replica:0/task:0/device:GPU:0

def step_fn(x):
  return x * 2

result = 0
for i in range(10):
  result += strategy.run(step_fn, args=(i,))
print(result)  # 90

os.environ['TF_CONFIG'] = json.dumps({
  'cluster': {
      'worker': ["localhost:12345", "localhost:23456"],
      'ps': ["localhost:34567"]
  },
  'task': {'type': 'worker', 'index': 0}
})

# This implicitly uses TF_CONFIG for the cluster and current task info.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

...

if strategy.cluster_resolver.task_type == 'worker':
  # Perform something that's only applicable on workers. Since we set this
  # as a worker above, this block will run on this particular instance.
elif strategy.cluster_resolver.task_type == 'ps':
  # Perform something that's only applicable on parameter servers. Since we
  # set this as a worker above, this block will not run on this particular
  # instance.

Методы

distribute_datasets_from_function

Просмотреть исходный код

distribute_datasets_from_function(
    dataset_fn, options=None
)

Распределяет экземпляры tf.data.Dataset, созданные вызовами dataset_fn.

dataset_fn будет вызываться для каждого рабочего узла стратегии. В этом случае у нас только один рабочий узел и одно устройство, поэтому dataset_fn будет вызвана один раз.

dataset_fn должна принимать экземпляр tf.distribute.InputContext, через который можно получить сведения о пакетировании и репликации входных данных:

def dataset_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  d = tf.data.Dataset.from_tensors([[1.]]).repeat().batch(batch_size)
  return d.shard(
      input_context.num_input_pipelines, input_context.input_pipeline_id)

inputs = strategy.distribute_datasets_from_function(dataset_fn)

for batch in inputs:
  replica_results = strategy.run(replica_fn, args=(batch,))

experimental_distribute_dataset

Просмотреть исходный код

experimental_distribute_dataset(
    dataset, options=None
)

Распределяет экземпляр tf.data.Dataset, переданный через dataset.

В этом случае используется только одно устройство, поэтому метод лишь оборачивает входной набор данных. Однако входные данные будут предварительно загружены на указанное устройство. По возвращенному распределенному набору данных можно выполнять итерацию так же, как по обычному набору данных.

Примечание: В настоящее время пользователь не может добавлять дополнительные преобразования к распределенному набору данных.

Пример:

strategy = tf.distribute.OneDeviceStrategy()
dataset = tf.data.Dataset.range(10).batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)
for x in dist_dataset:
  print(x)  # [0, 1], [2, 3],...

Аргументы: dataset: tf.data.Dataset, который требуется предварительно загрузить на устройство. options: tf.distribute.InputOptions, используемый для управления параметрами распределения этого набора данных. Возвращает: «распределенный Dataset», по которому вызывающий код может выполнять итерацию.

experimental_distribute_values_from_function

Просмотреть исходный код

experimental_distribute_values_from_function(
    value_fn
)

Создает tf.distribute.DistributedValues из value_fn.

Эта функция создает tf.distribute.DistributedValues для передачи в run, reduce или другие методы, принимающие распределенные значения, если наборы данных не используются.

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

1. Возврат постоянного значения для каждой реплики:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   def value_fn(ctx):
     return tf.constant(1.)
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
          value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (<tf.Tensor: shape=(), dtype=float32, numpy=1.0>,
       <tf.Tensor: shape=(), dtype=float32, numpy=1.0>)
       

2. Распределение значений массива на основе replica_id:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   array_value = np.array([3., 2., 1.])
   def value_fn(ctx):
     return array_value[ctx.replica_id_in_sync_group]
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
           value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (3.0, 2.0)
       

3. Задание значений с помощью num_replicas_in_sync:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   def value_fn(ctx):
     return ctx.num_replicas_in_sync
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
           value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (2, 2)
       

4. Размещение значений на устройствах и их распределение:

   strategy = tf.distribute.TPUStrategy()
   worker_devices = strategy.extended.worker_devices
   multiple_values = []
   for i in range(strategy.num_replicas_in_sync):
     with tf.device(worker_devices[i]):
       multiple_values.append(tf.constant(1.0))
   
   def value_fn(ctx):
     return multiple_values[ctx.replica_id_in_sync_group]
   
   distributed_values = strategy.
     experimental_distribute_values_from_function(
     value_fn)
   

experimental_local_results

Просмотреть исходный код

experimental_local_results(
    value
)

Возвращает список всех локальных значений для реплик, содержащихся в value.

В OneDeviceStrategy предполагается, что value всегда является одним значением, поэтому результатом будет кортеж, содержащий это значение.

gather

Просмотреть исходный код

gather(
    value, axis
)

Собирает value с реплик по измерению axis на текущее устройство.

Для объекта tf.distribute.DistributedValues или объекта, подобного tf.Tensor, value этот API собирает и объединяет конкатенацией value с реплик по измерению axis. Результат копируется на «текущее» устройство, которым обычно является процессор рабочего узла, на котором выполняется программа. Для tf.distribute.TPUStrategy это первый хост TPU. Для многоклиентской tf.distribute.MultiWorkerMirroredStrategy это процессор каждого рабочего узла.

Этот API можно вызывать только в контексте межрепличных операций. Аналог в контексте реплики см. в разделе tf.distribute.ReplicaContext.all_gather.

Примечание: Для всех стратегий, кроме tf.distribute.TPUStrategy, входной value на разных репликах должен иметь одинаковый ранг, а его формы должны совпадать по всем измерениям, кроме измерения axis. Иными словами, формы не могут различаться по измерению d, где d не равно аргументу axis. Например, если задан объект tf.distribute.DistributedValues с компонентными тензорами формы (1, 2, 3) и (1, 3, 3) на двух репликах, для него можно вызвать gather(..., axis=1, ...), но нельзя вызвать gather(..., axis=0, ...) или gather(..., axis=2, ...). Однако для tf.distribute.TPUStrategy.gather все тензоры должны иметь одинаковый ранг и форму.

Примечание: Компонентные тензоры объекта tf.distribute.DistributedValues value должны иметь ненулевой ранг. В противном случае перед сбором рассмотрите возможность использования tf.expand_dims.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# A DistributedValues with component tensor of shape (2, 1) on each replica
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(tf.constant([[1], [2]])))
@tf.function
def run():
  return strategy.gather(distributed_values, axis=0)
run()
<tf.Tensor: shape=(4, 1), dtype=int32, numpy=
array([[1],
       [2],
       [1],
       [2]], dtype=int32)>

Дополнительные сочетания параметров показаны в следующем примере:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1", "GPU:2", "GPU:3"])
single_tensor = tf.reshape(tf.range(6), shape=(1,2,3))
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(single_tensor))
@tf.function
def run(axis):
  return strategy.gather(distributed_values, axis=axis)
axis=0
run(axis)
<tf.Tensor: shape=(4, 2, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=1
run(axis)
<tf.Tensor: shape=(1, 8, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=2
run(axis)
<tf.Tensor: shape=(1, 2, 12), dtype=int32, numpy=
array([[[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5]]], dtype=int32)>

reduce

Просмотреть исходный код

reduce(
    reduce_op, value, axis
)

Выполняет сведение value по репликам.

В OneDeviceStrategy есть только одна реплика, поэтому при axis=None значение просто возвращается. Если для axis указано значение, отличное от None, например axis=0, значение сводится по этой оси и возвращается.

Пример:

t = tf.range(10)

result = strategy.reduce(tf.distribute.ReduceOp.SUM, t, axis=None).numpy()
# result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

result = strategy.reduce(tf.distribute.ReduceOp.SUM, t, axis=0).numpy()
# result: 45

run

Просмотреть исходный код

run(
    fn, args=(), kwargs=None, options=None
)

Запускает fn на каждой реплике с указанными аргументами.

В OneDeviceStrategy метод fn просто вызывается в области видимости устройства для указанного устройства с переданными аргументами.

scope

Просмотреть исходный код

scope()

Возвращает менеджер контекста, устанавливающий эту стратегию как текущую.

Внутри блока кода with strategy.scope(): этот поток использует создатель переменных, заданный методом strategy, и входит в «межрепличный контекст».

В OneDeviceStrategy все переменные, созданные внутри strategy.scope(), будут размещены на device, указанном при создании стратегии. Пример см. в документации этого класса.