tf.compat.v2.distribute.experimental.CentralStorageStrategy

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

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

tf.compat.v2.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.experimental_run_v2(train_step, args=(x,))

Методы

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset
)

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

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.experimental_run_v2(replica_fn, args=(batch,))

experimental_local_results

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

experimental_local_results(
    value
)

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

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

experimental_make_numpy_dataset

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

experimental_make_numpy_dataset(
    numpy_input
)

Создаёт tf.data.Dataset для входных данных, предоставленных через массив NumPy.

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

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

Пример:

numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_run_v2

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

experimental_run_v2(
    fn, args=(), kwargs=None
)

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

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

reduce

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

reduce(
    reduce_op, value, axis
)

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

Учитывая значение для каждой реплики, возвращаемое experimental_run_v2, например, потерю на пример, пакет будет разделён между всеми репликами. Эта функция позволяет агрегировать по репликам и, необязательно, по элементам пакета. Например, если у вас есть глобальный размер пакета 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.experimental_run_v2(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

scope

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

scope()

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

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