tf.distribute.experimental.CentralStorageStrategy

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

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

tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=None, parameter_device=None
)

Переменные назначаются локальному процессору или единственному графическому процессору. Если доступно более одного графического процессора, вычисления (кроме операций обновления переменных) будут дублированы на всех графических процессорах.

Пример:

strategy = tf.distribute.experimental.CentralStorageStrategy()
# Create a dataset
ds = tf.data.Dataset.range(5).batch(2)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(ds)

with strategy.scope():
  @tf.function
  def train_step(val):
    return val + 1

  # Iterate over the distributed dataset
  for x in dist_dataset:
    # process dataset elements
    strategy.run(train_step, args=(x,))

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 — это функция-вход, которая имеет аргумент tf.distribute.InputContext и возвращает экземпляр tf.data.Dataset. Ожидается, что возвращаемый из dataset_fn набор данных уже разбит на пакеты для каждой реплики (т. е. общий размер пакета разделен на количество реплик в синхронизации) и разделен. tf.distribute.Strategy.distribute_datasets_from_function не разбивает и не разделяет экземпляр tf.data.Dataset, возвращаемый функцией ввода. dataset_fn будет вызвана на процессорном узле каждого из рабочих узлов, и каждый создает набор данных, где каждая реплика на этом рабочем узле будет извлекать один пакет входных данных (т. е. если на рабочем узле две реплики, из Dataset каждый шаг будут извлечены два пакета).

Этот метод может использоваться для нескольких целей. Во-первых, он позволяет указать свою собственную логику разбиения на пакеты и фрагментации. (В отличие от tf.distribute.experimental_distribute_dataset, который выполняет разбиение на пакеты и фрагментацию за вас.) Например, где experimental_distribute_dataset не может разбить входные файлы, этот метод может использоваться для ручного разбиения набора данных (избегая медленного поведения по умолчанию в experimental_distribute_dataset). В случаях, когда набор данных бесконечен, это разбиение можно выполнить, создав реплики наборов данных, отличающиеся только своим случайным начальным значением.

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

Вы можете использовать свойство element_spec возвращаемого этим API tf.distribute.DistributedDataset, чтобы запросить tf.TypeSpec элементов, возвращаемых итератором. Это можно использовать для установки свойства input_signature tf.function. Следуйте tf.distribute.DistributedDataset.element_spec, чтобы увидеть пример.

Примечание: Если вы используете TPUStrategy, порядок обработки данных рабочими узлами при использовании tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function не гарантируется. Это обычно необходимо, если вы используете tf.distribute для масштабирования прогнозирования. Однако вы можете вставить индекс для каждого элемента в пакете и упорядочить выходы соответствующим образом. Обратитесь к этому фрагменту за примером того, как упорядочить выходы.

Примечание: Состояние наборов данных преобразования в настоящее время не поддерживаются с tf.distribute.experimental_distribute_dataset или tf.distribute.distribute_datasets_from_function. Любые операции с состоянием, которые может иметь набор данных, в настоящее время игнорируются. Например, если ваш набор данных имеет map_fn использующий tf.random.uniform для поворота изображения, то у вас есть граф набора данных, зависящий от состояния (т. е. случайного начального значения) на локальной машине, где выполняется процесс python.

Для получения дополнительных сведений и свойств этого метода обратитесь к учебнику по распределенному вводу. Если вас интересует обработка последнего частичного пакета, прочитайте эту секцию.

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset, options=None
)

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

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

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

Пример:

strategy = tf.distribute.CentralStorageStrategy()  # with 1 CPU and 1 GPU
dataset = tf.data.Dataset.range(10).batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)
for x in dist_dataset:
  print(x)  # Prints PerReplica values [0, 1], [2, 3],...

Args: dataset: tf.data.Dataset, который должен быть предварительно загружен на устройство. options: tf.distribute.InputOptions для управления параметрами распределения набора данных.

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

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

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)

1. Указание значений с помощью 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)

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.

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

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 по репликам.

Учитывая значение на реплике, возвращенное run, например, потерю на пример, пакет будет разделен между всеми репликами. Эта функция позволяет агрегировать по репликам и, по желанию, также по элементам пакета. Например, если у вас есть глобальный размер пакета 8 и 2 реплики, значения для примеров [0, 1, 2, 3] будут на реплике 0, а [4, 5, 6, 7] — на реплике 1. По умолчанию reduce просто агрегирует по репликам, возвращая [0+4, 1+5, 2+6, 3+7]. Это полезно, когда каждая реплика вычисляет скаляр или какое-то другое значение, у которого нет измерения "пакет" (например, градиент). Чаще всего вы захотите агрегировать по глобальному пакету, что можно сделать, указав измерение пакета как axis, обычно axis=0. В этом случае он вернет скаляр 0+1+2+3+4+5+6+7.

Если есть последний частичный пакет, вам нужно указать ось, чтобы размерность результата была согласованной между репликами. Так, если последний пакет имеет размер 6 и он разделен на [0, 1, 2, 3] и [4, 5], у вас возникнет несоответствие размерностей, если вы не укажете axis=0. Если вы укажете tf.distribute.ReduceOp.MEAN, используя axis=0 будет использоваться правитель 6. Противопоставьте это вычислению reduce_mean для получения скалярного значения на каждой реплике и этой функции для усреднения этих средних значений, что будет учитывать некоторые значения 1/8 и другие 1/4.

Пример:

strategy = tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=['CPU:0', 'GPU:0'], parameter_device='CPU:0')
ds = tf.data.Dataset.range(10)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(ds)

with strategy.scope():
  @tf.function
  def train_step(val):
    # pass through
    return val

  # Iterate over the distributed dataset
  for x in dist_dataset:
    result = strategy.run(train_step, args=(x,))

result = strategy.reduce(tf.distribute.ReduceOp.SUM, result,
                         axis=None).numpy()
# result: array([ 4,  6,  8, 10])

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

run

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

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

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

В CentralStorageStrategy, fn вызывается на каждой вычислительной реплике с предоставленными аргументами "на реплику", специфичными для этого устройства.

scope

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

scope()

Управляющая конструкция для установки стратегии в текущий статус и распределения переменных.

Этот метод возвращает управляющую конструкцию и используется следующим образом:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# Variable created inside scope:
with strategy.scope():
  mirrored_variable = tf.Variable(1.)
mirrored_variable
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>,
  1: <tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=1.0>
}
# Variable created outside scope:
regular_variable = tf.Variable(1.)
regular_variable
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

Что происходит при входе в область действия Strategy.scope?

• strategy устанавливается в глобальном контексте как текущая стратегия. Внутри этой области, tf.distribute.get_strategy() теперь вернет эту стратегию. За пределами этой области, она возвращает стратегию по умолчанию no-op.
• Вход в область также означает вход в "контекст между репликами". См. tf.distribute.StrategyExtended для объяснения контекстов между репликами и реплики.
• Создание переменной внутри scope перехватывается стратегией. Каждая стратегия определяет, как она хочет повлиять на создание переменной. Синхронные стратегии, такие как MirroredStrategy, TPUStrategy и MultiWorkerMiroredStrategy, создают переменные, дублированные на каждой реплике, в то время как ParameterServerStrategy создает переменные на параметрических серверах. Это делается с помощью пользовательского tf.variable_creator_scope.
• В некоторых стратегиях также может быть введена область действия по умолчанию: в MultiWorkerMiroredStrategy, область действия устройства по умолчанию "/CPU:0" вводится на каждом узле.

Примечание: Вход в область действия не автоматически распределяет вычисление, за исключением случая высокоуровневых фреймворков обучения, таких как keras model.fit. Если вы не используете model.fit, вам нужно использовать API strategy.run для явного распределения этого вычисления. См. пример в учебнике по кастомному циклу обучения .

Что должно быть в области и что вне?

Существует ряд требований к тому, что должно происходить внутри области. Однако в тех местах, где у нас есть информация о используемой стратегии, мы часто входим в область для пользователя, чтобы ему не приходилось делать это явно (т.е. вызов внутри или вне области допустим).

• Все, что создает переменные, которые должны быть распределенными, должно вызываться в strategy.scope. Это можно сделать, либо напрямую вызвав функцию создания переменной в контексте области видимости, либо, полагаясь на другой API, например, strategy.run или keras.Model.fit, для автоматического включения в него. Любые переменные, созданные вне области видимости, не будут распределены и могут повлиять на производительность. Некоторые распространенные объекты, которые создают переменные в TF, это Models, Optimizers, Metrics. Такие объекты всегда должны инициализироваться в области видимости, и любые функции, которые могут создавать переменные лениво (например, Model.__call__(), отслеживание tf.function и т.д.), должны аналогично вызываться в пределах области видимости. Другим источником создания переменных может быть восстановление из контрольной точки — когда переменные создаются лениво. Обратите внимание, что любая переменная, созданная внутри стратегии, сохраняет информацию о стратегии. Таким образом, чтение и запись этих переменных вне strategy.scope также могут работать безупречно, без необходимости для пользователя входить в область видимости.
• Некоторые API стратегий (например, strategy.run и strategy.reduce), которые требуют нахождения в области видимости стратегии, автоматически входят в область видимости, что означает, что при использовании этих API вам не нужно явно входить в область видимости.
• Когда tf.keras.Model создается внутри strategy.scope, объект Model сохраняет информацию о области видимости. При последующем вызове методов высокоуровневого фреймворка обучения, таких как model.compile, model.fit и т.д., захваченная область видимости будет автоматически введена, и соответствующая стратегия будет использована для распределения обучения и т.д. Подробный пример см. в руководстве по распределенному Keras. ПРЕДУПРЕЖДЕНИЕ: простой вызов model(..) не автоматически входит в захваченную область видимости — только высокоуровневые API фреймворка обучения поддерживают это поведение: model.compile, model.fit, model.evaluate, model.predict и model.save могут вызываться внутри или вне области видимости.
• Следующее может находиться как внутри, так и вне области видимости:
  • Создание входных наборов данных
  • Определение tf.function, представляющих ваш шаг обучения
  • API сохранения, такие как tf.saved_model.save. Загрузка создает переменные, поэтому она должна находиться внутри области видимости, если вы хотите обучить модель распределённым способом.
  • Сохранение контрольных точек. Как упоминалось выше — checkpoint.restore иногда может потребоваться внутри области видимости, если оно создаёт переменные.