tf.compat.v2.distribute.OneDeviceStrategy

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

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

tf.compat.v2.distribute.OneDeviceStrategy(
    device
)

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

Типичное использование этой стратегии — тестирование вашего кода с помощью 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.experimental_run_v2(step_fn, args=(i,))
print(result)  # 90

Методы

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

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 с возвращаемым набором данных для дальнейшего распределения его с помощью стратегии.

Пример:

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

В OneDeviceStrategy, fn просто вызывается в области устройства для данного устройства с предоставленными аргументами.

reduce

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

reduce(
    reduce_op, value, axis
)

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

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

scope

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

scope()

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

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

В OneDeviceStrategy, все переменные, созданные внутри strategy.scope(), будут на device, указанном во время создания стратегии. См. пример в документации по этому классу.