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 объекта tf.distribute.DistributedDataset, возвращаемого этим API, можно использовать для получения 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],...

Аргументы: 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>)
       

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.

В 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, а реплик две, значения для примеров [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() вернёт эту стратегию. За пределами этой области он возвращает стратегию по умолчанию, не выполняющую никаких действий.
• При входе в область также устанавливается «контекст между репликами». Объяснение контекстов между репликами и контекстов реплик см. в разделе 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, создающие переменные: модели, оптимизаторы и метрики. Такие объекты всегда следует инициализировать в области действия; функции, которые могут создавать переменные с задержкой (например, 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 необходимо выполнять внутри области действия, если при этом создаются переменные.