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.

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

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

Пример:

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. Распределение значений в массиве в зависимости от идентификатора реплики:

   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, указанном во время создания стратегии. См. пример в документации по этому классу.