tf.distribute.OneDeviceStrategy

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

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

tf.distribute.OneDeviceStrategy(
    device
)

Использование этой стратегии разместит все переменные, созданные в её области, на указанном устройстве. Входные данные, распределённые с помощью этой стратегии, будут предварительно загружены на указанное устройство. Кроме того, любые функции, вызываемые через strategy.experimental_run_v2 также будут размещены на указанном устройстве.

Типичное использование этой стратегии может заключаться в тестировании вашего кода с помощью API tf.distribute.Strategy перед переключением на другие стратегии, которые фактически распределяют вычисления по нескольким устройствам/машинам.

Например:

tf.enable_eager_execution()
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.experimental_run_v2(step_fn, args=(i,))
print(result)  # 90

Методы

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset
)

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

Возвращаемый распределённый набор данных можно итерировать аналогично тому, как можно итерировать обычные наборы данных. ПРИМЕЧАНИЕ: В настоящее время пользователь не может добавлять дополнительные преобразования к распределённому набору данных.

Следующий пример:

strategy = tf.distribute.MirroredStrategy()

# Create a dataset
dataset = dataset_ops.Dataset.TFRecordDataset([
  "/a/1.tfr", "/a/2.tfr", "/a/3.tfr", "/a/4.tfr"])

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

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

В многоузловой среде мы сначала попробуем распределить набор данных, пытаясь определить, создается ли набор данных из ReaderDatasets (например, TFRecordDataset, TextLineDataset и т. д.), и, если это так, попробуем разбить входные файлы. Обратите внимание, что должно быть по крайней мере один входной файл на каждый рабочий узел. Если у вас меньше одного входного файла на рабочий узел, рекомендуется отключить распределение вашего набора данных с помощью метода ниже.

Если эта попытка окажется неудачной (например, набор данных создан из Dataset.range), мы разделим набор данных равномерно в конце, добавив операцию .shard в конец потока обработки. Это приведет к тому, что весь процесс подготовки данных для всех данных будет выполнен на каждом рабочем узле, и каждый рабочий узел будет выполнять избыточную работу. Мы выведем предупреждение, если этот метод фрагментации будет выбран. В этом случае рассмотрите возможность использования experimental_distribute_datasets_from_function вместо этого.

Вы можете отключить фрагментацию наборов данных между рабочими узлами, используя опцию auto_shard в tf.data.experimental.DistributeOptions.

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

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

experimental_distribute_datasets_from_function

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

experimental_distribute_datasets_from_function(
    dataset_fn
)

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

dataset_fn будет вызываться один раз для каждого рабочего узла в стратегии. Каждая реплика на этом рабочем узле будет декьюить одну партию входных данных из локального Dataset (т. е. если у рабочего узла есть две реплики, две партии будут декьюиться из Dataset на каждом шаге).

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

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.

Примечание: Это возвращает только значения на рабочем узле, инициированном этим клиентом. При использовании tf.distribute.Strategy, такого как tf.distribute.experimental.MultiWorkerMirroredStrategy, каждый рабочий узел будет своим клиентом, и эта функция вернёт только значения, вычисленные на этом рабочем узле.

experimental_make_numpy_dataset

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

experimental_make_numpy_dataset(
    numpy_input, session=None
)

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

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

Обратите внимание, что, скорее всего, вам потребуется использовать tf.distribute.Strategy.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

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

experimental_run(
    fn, input_iterator=None
)

Выполняет операции в fn на каждой реплике с входными данными из input_iterator.

УСТАРЕЛО: Этот метод недоступен в TF 2.x. Пожалуйста, переключитесь на использование experimental_run_v2 вместо этого.

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

Каждая реплика получит один отдельный вход из входных данных, предоставленных одним вызовом get_next на итераторе входных данных.

fn может вызвать tf.distribute.get_replica_context(), чтобы получить доступ к членам, таким как replica_id_in_sync_group.

experimental_run_v2

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

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

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

Выполняет операции, определённые fn на каждой реплике. Если args или kwargs содержат значения "по реплике", такие как те, что получены с помощью "распределённого Dataset", при выполнении fn на конкретной реплике, она будет выполнена с компонентом этих значений "по реплике", соответствующим этой реплике.

fn может вызвать tf.distribute.get_replica_context(), чтобы получить доступ к членам, таким как all_reduce.

Все аргументы в args или kwargs должны быть либо вложенными тензорами, либо объектами "по реплике", содержащими тензоры или составные тензоры.

make_dataset_iterator

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

make_dataset_iterator(
    dataset
)

Создаёт итератор для входных данных, предоставленных через dataset.

УСТАРЕЛО: Этот метод недоступен в TF 2.x.

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

Пользователь также может использовать make_input_fn_iterator , если хочет настроить, какой ввод подаётся на какую реплику/рабочий процесс и т. д.

make_input_fn_iterator

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

make_input_fn_iterator(
    input_fn, replication_mode=tf.distribute.InputReplicationMode.PER_WORKER
)

Возвращает итератор, разделённый между репликами, созданный из функции ввода.

УСТАРЕЛО: Этот метод недоступен в TF 2.x.

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

def input_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)
with strategy.scope():
  iterator = strategy.make_input_fn_iterator(input_fn)
  replica_results = strategy.experimental_run(replica_fn, iterator)

Возвращаемый tf.data.Dataset из input_fn должен иметь размер пакета на реплику, который можно вычислить с помощью input_context.get_per_replica_batch_size.

reduce

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

reduce(
    reduce_op, value, axis=None
)

Сведение 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.

scope

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

scope()

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

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

update_config_proto

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

update_config_proto(
    config_proto
)

Возвращает копию config_proto , изменённую для использования с этой стратегией.

УСТАРЕЛО: Этот метод недоступен в TF 2.x.

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