tf.compat.v1.distribute.Strategy

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

tf.compat.v1.distribute.Strategy(
    extended
)

См. руководство для обзора и примеров.

Примечание: Не все tf.distribute.Strategy реализации в настоящее время поддерживают разбиение переменных TensorFlow (где одна переменная разделена между несколькими устройствами).

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 будет вызываться на устройстве CPU каждого из узлов, и каждый из них сгенерирует набор данных, где каждая реплика на этом узле будет извлекать одну партию входных данных (т. е. если у узла две реплики, две партии будут извлекаться из 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.distribute.DistributedDataset из tf.data.Dataset.

Возвращаемый tf.distribute.DistributedDataset можно перебирать, как обычные наборы данных. ПРИМЕЧАНИЕ: пользователь не может добавить больше преобразований к tf.distribute.DistributedDataset. Вы можете только создать итератор или проверить tf.TypeSpec генерируемых им данных. См. документацию API tf.distribute.DistributedDataset для получения дополнительной информации.

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

global_batch_size = 2
# Passing the devices is optional.
strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
# Create a dataset
dataset = tf.data.Dataset.range(4).batch(global_batch_size)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)
@tf.function
def replica_fn(input):
  return input*2
result = []
# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  result.append(strategy.run(replica_fn, args=(x,)))
print(result)
[PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([0])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([2])>
}, PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([4])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([6])>
}]

Три ключевых действия, происходящие внутри этого метода: группировка, разбиение и предварительная выборка.

В фрагменте кода выше, dataset группируется по global_batch_size, и вызов experimental_distribute_dataset по нему перегруппировывает dataset до нового размера пакета, равного глобальному размеру пакета, деленному на количество реплик в синхронизации. Мы перебираем его с помощью Python-цикла. x — это tf.distribute.DistributedValues, содержащий данные для всех реплик, и каждая реплика получает данные нового размера пакета. tf.distribute.Strategy.run позаботится о подаче правильных данных на реплику в x нужной replica_fn операции, выполняемой на каждой реплике.

Разбиение включает автоматическое разбиение по нескольким узлам и внутри каждого узла. Во-первых, при распределённом обучении по нескольким узлам (т. е. при использовании tf.distribute.experimental.MultiWorkerMirroredStrategy или tf.distribute.TPUStrategy) автоматическое разбиение набора данных по нескольким узлам означает, что каждому узлу назначается подмножество всего набора данных (если установлен соответствующий tf.data.experimental.AutoShardPolicy). Это гарантирует, что на каждом шаге каждый узел будет обрабатывать глобальный пакет неперекрывающихся элементов набора данных. Автоматическое разбиение имеет несколько различных вариантов, которые можно задать с помощью tf.data.experimental.DistributeOptions. Затем разбиение внутри каждого узла означает, что метод разделит данные между всеми устройствами узла (если их больше одного).

Примечание: для автофрагментации по нескольким рабочим процессам, режим по умолчанию — tf.data.experimental.AutoShardPolicy.AUTO. Этот режим попытается разбить входной набор данных по файлам, если набор данных создается из наборов данных читателя (например, tf.data.TFRecordDataset, tf.data.TextLineDataset и т. д.) или в противном случае разбить набор данных по данным, где каждый из рабочих процессов будет читать весь набор данных и обрабатывать только назначенный ему фрагмент. Однако, если у вас менее одного входного файла на рабочий процесс, мы рекомендуем отключить автофрагментацию набора данных по рабочим процессам, установив tf.data.experimental.DistributeOptions.auto_shard_policy в значение tf.data.experimental.AutoShardPolicy.OFF.

По умолчанию этот метод добавляет преобразование prefetch в конце предоставленного пользователем экземпляра tf.data.Dataset. Аргумент преобразования prefetch, который равен buffer_size, равен количеству реплик в синхронизации.

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

Примечание: Если вы используете 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_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. Пожалуйста, переключитесь на использование run вместо этого.

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

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

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

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 по репликам и вернуть результат на текущем устройстве.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
total = strategy.reduce("SUM", per_replica_result, axis=None)
total
<tf.Tensor: shape=(), dtype=int32, numpy=1>

Чтобы увидеть, как это будет выглядеть с несколькими репликами, рассмотрите тот же пример с MirroredStrategy с 2 графическими процессорами:

strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
# Check devices on which per replica result is:
strategy.experimental_local_results(per_replica_result)[0].device
# /job:localhost/replica:0/task:0/device:GPU:0
strategy.experimental_local_results(per_replica_result)[1].device
# /job:localhost/replica:0/task:0/device:GPU:1

total = strategy.reduce("SUM", per_replica_result, axis=None)
# Check device on which reduced result is:
total.device
# /job:localhost/replica:0/task:0/device:CPU:0

Этот API обычно используется для агрегирования результатов, возвращаемых различными репликами, для отчётности и т. д. Например, потерю, вычисленную различными репликами, можно усреднить с помощью этого API перед выводом.

Примечание: Результат копируется на «текущее» устройство — обычно это процессор работника, на котором выполняется программа. Для TPUStrategy, это первый хост TPU. Для многоклиентского MultiWorkerMirroredStrategy, это процессор каждого работника.

Существует ряд различных API tf.distribute для сворачивания значений по репликам:

• tf.distribute.ReplicaContext.all_reduce: Это отличается от Strategy.reduce тем, что предназначен для контекста реплики и не копирует результаты на устройство хоста. all_reduce обычно используется для вычислений внутри этапа обучения, таких как градиенты.
• tf.distribute.StrategyExtended.reduce_to и tf.distribute.StrategyExtended.batch_reduce_to: Эти API являются более продвинутыми версиями Strategy.reduce, так как они позволяют настраивать место назначения результата. Они также вызываются в контексте между репликами.

Каким должен быть ось?

Учитывая значение на реплику, возвращаемое run, скажем, потерю на пример, пакет будет разделен между всеми репликами. Эта функция позволяет агрегировать по репликам, а также по элементам пакета, указав параметр оси соответственно.

Например, если у вас есть глобальный размер пакета 8 и 2 реплики, значения для примеров [0, 1, 2, 3] будут на реплике 0, а [4, 5, 6, 7] — на реплике 1. С axis=None, reduce будет агрегировать только по репликам, возвращая [0+4, 1+5, 2+6, 3+7]. Это полезно, когда каждая реплика вычисляет скаляр или другое значение без измерения «пакета» (например, градиент или потерю).

strategy.reduce("sum", per_replica_result, axis=None)

Иногда вам нужно будет агрегировать как по глобальному пакету, так и по всем репликам. Вы можете получить это поведение, указав размер пакета как ось, обычно axis=0. В этом случае он вернёт скаляр 0+1+2+3+4+5+6+7.

strategy.reduce("sum", per_replica_result, axis=0)

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

run

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

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

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

Этот метод является основным способом распределения вычислений с объектом tf.distribute. Он вызывает fn на каждой реплике. Если у args или kwargs есть tf.distribute.DistributedValues, такие как те, которые создаются tf.distribute.DistributedDataset из tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function, когда fn выполняется на определенной реплике, она выполняется с компонентом tf.distribute.DistributedValues, соответствующим этой реплике.

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

Все аргументы в args или kwargs должны быть либо значениями Python, либо вложенной структурой тензоров, например, списком тензоров, в этом случае args и kwargs будут переданы в fn , вызываемый на каждой реплике. Либо args или kwargs могут быть tf.distribute.DistributedValues, содержащими тензоры или составные тензоры, т. е. tf.compat.v1.TensorInfo.CompositeTensor, в этом случае каждый вызов fn получит компонент tf.distribute.DistributedValues, соответствующий его реплике.

Пример использования:

1. Входной тензор с постоянным значением.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
tensor_input = tf.constant(3.0)
@tf.function
def replica_fn(input):
  return input*2.0
result = strategy.run(replica_fn, args=(tensor_input,))
result
PerReplica:{
  0: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>,
  1: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>
}

1. Вход DistributedValues.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
@tf.function
def run():
  def value_fn(value_context):
    return value_context.num_replicas_in_sync
  distributed_values = (
    strategy.experimental_distribute_values_from_function(
      value_fn))
  def replica_fn2(input):
    return input*2
  return strategy.run(replica_fn2, args=(distributed_values,))
result = run()
result
<tf.Tensor: shape=(), dtype=int32, numpy=4>

1. Используйте tf.distribute.ReplicaContext для allreduce значений.

strategy = tf.distribute.MirroredStrategy(["gpu:0", "gpu:1"])
@tf.function
def run():
   def value_fn(value_context):
     return tf.constant(value_context.replica_id_in_sync_group)
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
           value_fn))
   def replica_fn(input):
     return tf.distribute.get_replica_context().all_reduce("sum", input)
   return strategy.run(replica_fn, args=(distributed_values,))
result = run()
result
PerReplica:{
  0: <tf.Tensor: shape=(), dtype=int32, numpy=1>,
  1: <tf.Tensor: shape=(), dtype=int32, numpy=1>
}

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 или model.fit, чтобы он вошёл в область за вас. Любая переменная, созданная вне области, не будет распределена и может повлиять на производительность. Общие вещи, которые создают переменные в TF: модели, оптимизаторы, метрики. Эти всегда должны создаваться внутри области. Ещё одним источником создания переменных может быть восстановление контрольной точки — когда переменные создаются лениво. Обратите внимание, что любая переменная, созданная внутри стратегии, сохраняет информацию о стратегии. Поэтому чтение и запись этих переменных за пределами strategy.scope также могут работать без проблем, без необходимости ввода пользователем области.
• Некоторые API стратегий (например, strategy.run и strategy.reduce), которые требуют наличия в области стратегии, автоматически вводят область, что означает, что при использовании этих API вам не нужно входить в область самостоятельно.
• Когда tf.keras.Model создаётся внутри strategy.scope, мы захватываем эту информацию. При вызове методов высокоуровневых фреймворков обучения, таких как 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 иногда может потребоваться быть внутри области, если она создаёт переменные.

update_config_proto

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

update_config_proto(
    config_proto
)

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

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

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