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.

Методы

experimental_assign_to_logical_device

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

experimental_assign_to_logical_device(
    tensor, logical_device_id
)

Добавляет аннотацию, что tensor будет назначено логическому устройству.

Примечание: Этот API в настоящее время поддерживается только в TPUStrategy. Это добавляет аннотацию к tensor, указывающую, что операции с tensor будут вызваны на логическом устройстве с идентификатором logical_device_id. При использовании параллелизма моделей, по умолчанию все операции размещаются на нулевом логическом устройстве.

# Initializing TPU system with 2 logical devices and 4 replicas.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=4)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)
iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  output = tf.add(inputs, inputs)

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output

strategy.run(step_fn, args=(next(iterator),))

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset
)

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

experimental_distribute_datasets_from_function

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

experimental_distribute_datasets_from_function(
    dataset_fn
)

Распределяет экземпляры 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.experimental_distribute_datasets_from_function(dataset_fn)

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

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()
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>,)

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

strategy = tf.distribute.MirroredStrategy()
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,)

1. Указание значений с использованием num_replicas_in_sync:

strategy = tf.distribute.MirroredStrategy()
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
(1,)

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

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

experimental_make_numpy_dataset

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

experimental_make_numpy_dataset(
    numpy_input
)

Создаёт tf.data.Dataset из массива NumPy. (устарело)

Это позволяет избежать добавления numpy_input в качестве большой константы в график и копирует данные на машину или машины, которые будут обрабатывать входные данные.

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

Пример:

strategy = tf.distribute.MirroredStrategy()
numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dataset
<TensorSliceDataset shapes: (), types: tf.float32>
dataset = dataset.batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_replicate_to_logical_devices

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

experimental_replicate_to_logical_devices(
    tensor
)

Добавляет аннотацию, что tensor будет дублирован на все логические устройства.

Примечание: Этот API пока поддерживается только в TPUStrategy. Это добавляет аннотацию к тензору tensor, указывающую, что операции с tensor будут вызываться на всех логических устройствах.

# Initializing TPU system with 2 logical devices and 4 replicas.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=4)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  images, labels = inputs
  images = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)

  # For loss calculation, all logical devices share the same logits
  # and labels.
  labels = strategy.experimental_replicate_to_logical_devices(labels)
  output = strategy.experimental_replicate_to_logical_devices(output)
  loss = loss_fn(labels, output)

  return loss

strategy.run(step_fn, args=(next(iterator),))

Аргументы: tensor: входной тензор для аннотации.

experimental_split_to_logical_devices

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

experimental_split_to_logical_devices(
    tensor, partition_dimensions
)

Добавляет аннотацию, что tensor будет разделен между логическими устройствами.

Примечание: Этот API пока поддерживается только в TPUStrategy. Это добавляет аннотацию к тензору tensor указывающую, что операции с tensor будут разделены между несколькими логическими устройствами. Тензор tensor будет разделен по измерениям, указанным в partition_dimensions. Измерения tensor должны быть кратны соответствующему значению в partition_dimensions.

Например, для системы с 8 логическими устройствами, если tensor это тензор изображения с формой (размер_пакета, ширина, высота, канал) и partition_dimensions это [1, 2, 4, 1], то tensor будет разделен на 2 по ширине и 4 по высоте, и значения разделенного тензора будут переданы на 8 логических устройств.

# Initializing TPU system with 8 logical devices and 1 replica.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 2, 2, 2],
    num_replicas=1)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  inputs = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)
  return output

strategy.run(step_fn, args=(next(iterator),))

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

reduce

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

reduce(
    reduce_op, value, axis
)

Сведение value по репликам.

В OneDeviceStrategy, есть только одна реплика, поэтому, если ось=None, возвращается просто значение. Если ось указана как что-то другое, например ось=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 , указанном во время создания стратегии. См. пример в документации для этого класса.