Сохраните дату! Google I / O возвращается 18-20 мая Зарегистрируйтесь сейчас
Эта страница переведена с помощью Cloud Translation API.
Switch to English

Распределенное обучение с TensorFlow

Посмотреть на TensorFlow.org Запускаем в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

Обзор

tf.distribute.Strategy - это TensorFlow API для распределения обучения между несколькими графическими процессорами, несколькими компьютерами или TPU. Используя этот API, вы можете распространять существующие модели и обучающий код с минимальными изменениями кода.

tf.distribute.Strategy была разработана с учетом следующих ключевых целей:

  • Простота использования и поддержка нескольких сегментов пользователей, включая исследователей, инженеров машинного обучения и т. Д.
  • Обеспечьте хорошую производительность прямо из коробки.
  • Легкое переключение между стратегиями.

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

В TensorFlow 2.x вы можете выполнять свои программы с нетерпением или в виде графика с помощью tf.function . tf.distribute.Strategy намеревается поддерживать оба этих режима выполнения, но лучше всего работает с tf.function . Активный режим рекомендуется только для целей отладки и не поддерживается TPUStrategy . Хотя в этом руководстве основное внимание уделяется обучению, этот API также можно использовать для распределения оценок и прогнозов на различных платформах.

Вы можете использовать tf.distribute.Strategy с очень небольшими изменениями в вашем коде, потому что мы изменили базовые компоненты TensorFlow, чтобы они учитывали стратегию. Сюда входят переменные, слои, модели, оптимизаторы, метрики, сводки и контрольные точки.

В этом руководстве мы объясняем различные типы стратегий и то, как вы можете использовать их в разных ситуациях. Чтобы узнать, как отлаживать проблемы с производительностью, см. Руководство по оптимизации производительности графического процессора TensorFlow .

# Import TensorFlow
import tensorflow as tf

Типы стратегий

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

  • Синхронное и асинхронное обучение: это два распространенных способа распределения обучения с параллелизмом данных. При синхронном обучении все рабочие обучаются синхронно на разных срезах входных данных и агрегируют градиенты на каждом этапе. При асинхронном обучении все рабочие независимо обучаются над входными данными и асинхронно обновляют переменные. Обычно обучение синхронизации поддерживается с помощью архитектуры all-reduce и async через серверную архитектуру параметров.
  • Аппаратная платформа: вы можете масштабировать свое обучение на несколько графических процессоров на одной машине или на нескольких машинах в сети (с 0 или более графическими процессорами на каждой) или на облачных TPU.

Для поддержки этих вариантов использования доступно шесть стратегий. В следующем разделе объясняется, какие из них поддерживаются в каких сценариях в TF. Вот краткий обзор:

API обучения ЗеркальныйСтратегия TPUStrategy MultiWorkerMirroredStrategy CentralStorageСтратегия ParameterServerStrategy
Keras API Поддерживается Поддерживается Поддерживается Экспериментальная поддержка Поддерживаемая запланированная публикация 2.4
Пользовательский цикл обучения Поддерживается Поддерживается Поддерживается Экспериментальная поддержка Экспериментальная поддержка
API оценщика Ограниченная поддержка Не поддерживается Ограниченная поддержка Ограниченная поддержка Ограниченная поддержка

ЗеркальныйСтратегия

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

Для передачи обновлений переменных между устройствами используются эффективные алгоритмы полного сокращения. All-reduce объединяет тензоры на всех устройствах, складывая их, и делает их доступными на каждом устройстве. Это объединенный алгоритм, который очень эффективен и может значительно снизить накладные расходы на синхронизацию. В зависимости от типа связи, доступной между устройствами, существует множество доступных алгоритмов и реализаций all-reduce. По умолчанию в качестве реализации полного сокращения используется NVIDIA NCCL. Вы можете выбрать один из нескольких других вариантов или написать свой собственный.

Вот простейший способ создания MirroredStrategy :

mirrored_strategy = tf.distribute.MirroredStrategy()
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

Это создаст экземпляр MirroredStrategy который будет использовать все графические процессоры, видимые для TensorFlow, и использовать NCCL в качестве связи между устройствами.

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

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
WARNING:tensorflow:Some requested devices in `tf.distribute.Strategy` are not visible to TensorFlow: /job:localhost/replica:0/task:0/device:GPU:0,/job:localhost/replica:0/task:0/device:GPU:1
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')

Если вы хотите переопределить обмен данными между устройствами, вы можете сделать это с cross_device_ops аргумента cross_device_ops , предоставив экземпляр tf.distribute.CrossDeviceOps . В настоящее время tf.distribute.HierarchicalCopyAllReduce и tf.distribute.ReductionToOneDevice - это два параметра, отличные от tf.distribute.NcclAllReduce который используется по умолчанию.

mirrored_strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)

TPUStrategy

tf.distribute.TPUStrategy позволяет запускать обучение tf.distribute.TPUStrategy обработки тензорных (TPU). TPU - это специализированные ASIC от Google, предназначенные для значительного ускорения рабочих нагрузок машинного обучения. Они доступны в Google Colab, TensorFlow Research Cloud и Cloud TPU .

С точки зрения архитектуры распределенного обучения TPUStrategy - это та же MirroredStrategy - она ​​реализует синхронное распределенное обучение. TPU обеспечивают собственную реализацию эффективных операций all-reduce и других коллективных операций на нескольких ядрах TPU, которые используются в TPUStrategy .

Вот как можно создать экземпляр TPUStrategy :

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_address)
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
tpu_strategy = tf.distribute.TPUStrategy(cluster_resolver)

Экземпляр TPUClusterResolver помогает найти TPU. В Colab вам не нужно указывать для него какие-либо аргументы.

Если вы хотите использовать это для облачных TPU:

  • Вы должны указать имя вашего ресурса TPU в аргументе tpu .
  • Вы должны явно инициализировать систему TPU в начале программы. Это необходимо до того, как TPU можно будет использовать для вычислений. Инициализация системы TPU также стирает память TPU, поэтому важно сначала выполнить этот шаг, чтобы избежать потери состояния.

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy очень похож на MirroredStrategy . Он реализует синхронное распределенное обучение для нескольких рабочих, каждый из которых может иметь несколько графических процессоров. Подобно tf.distribute.MirroredStrategy , он создает копии всех переменных в модели на каждом устройстве для всех рабочих процессов.

Вот простейший способ создания MultiWorkerMirroredStrategy :

strategy = tf.distribute.MultiWorkerMirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',)
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.AUTO

MultiWorkerMirroredStrategy имеет две реализации для обмена данными между устройствами. CommunicationImplementation.RING основан на RPC и поддерживает как CPU, так и GPU. CommunicationImplementation.NCCL использует NCCL от Nvidia и обеспечивает высочайшую производительность на GPU, но не поддерживает CPU. CollectiveCommunication.AUTO оставляет выбор за Tensorflow. Вы можете указать их следующим образом:

communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options)
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',)
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.NCCL

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

ParameterServerStrategy

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

В обучении сервера параметров TensorFlow 2 используется архитектура на основе центрального координатора через класс tf.distribute.experimental.coordinator.ClusterCoordinator .

В этой реализации worker задачи и задачи parameter server запускают tf.distribute.Server которые прослушивают задачи от координатора. Координатор создает ресурсы, отправляет обучающие задачи, записывает контрольные точки и устраняет сбои задач.

При программировании, выполняемом на координаторе, вы будете использовать объект ParameterServerStrategy для определения шага обучения и использовать ClusterCoordinator для отправки шагов обучения удаленным работникам. Вот самый простой способ их создания:

strategy = tf.distribute.experimental.ParameterServerStrategy(
    tf.distribute.cluster_resolver.TFConfigClusterResolver(),
    variable_partitioner=variable_partitioner)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy)

Обратите внимание, что вам нужно будет настроить переменную среды TF_CONFIG, если вы используете TFConfigClusterResolver . Он похож на TF_CONFIG в MultiWorkerMirroredStrategy но имеет дополнительные предостережения.

В TF 1 ParameterServerStrategy доступен только с оценщиком через символ tf.compat.v1.distribute.experimental.ParameterServerStrategy .

CentralStorageСтратегия

tf.distribute.experimental.CentralStorageStrategy выполняет синхронное обучение. Переменные не зеркалируются, вместо этого они размещаются на ЦП, а операции реплицируются на всех локальных графических процессорах. Если есть только один графический процессор, все переменные и операции будут размещены на этом графическом процессоре.

Создайте экземпляр CentralStorageStrategy :

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()
INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:GPU:0'], variable_device = '/job:localhost/replica:0/task:0/device:GPU:0'

Это создаст экземпляр CentralStorageStrategy который будет использовать все видимые графические процессоры и ЦП. Обновление переменных на репликах будет агрегировано перед применением к переменным.

Другие стратегии

В дополнение к вышеуказанным стратегиям есть две другие стратегии, которые могут быть полезны для tf.distribute прототипов и отладки при использовании API tf.distribute .

Стратегия по умолчанию

Стратегия по умолчанию - это стратегия распространения, которая присутствует, когда в области нет явной стратегии распространения. Он реализует интерфейс tf.distribute.Strategy но является сквозным и не обеспечивает фактического распространения. Например, strategy.run(fn) просто вызовет fn . Код, написанный с использованием этой стратегии, должен вести себя точно так же, как код, написанный без какой-либо стратегии. Вы можете думать об этом как о «запретной» стратегии.

Стратегия по умолчанию - синглтон, и больше ее экземпляров создать нельзя. Его можно получить с помощью tf.distribute.get_strategy() вне любой явной области действия стратегии (тот же API, который можно использовать для получения текущей стратегии внутри области явной стратегии).

default_strategy = tf.distribute.get_strategy()

Эта стратегия служит двум основным целям:

  • Это позволяет безоговорочно писать код библиотеки с учетом распространения. Например, в tf.optimizer s можно использовать tf.distribute.get_strategy() и использовать эту стратегию для уменьшения градиентов - он всегда будет возвращать объект стратегии, для которого мы можем вызвать API сокращения.
# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None)  # reduce some values
1.0
  • Подобно библиотечному коду, его можно использовать для написания программ конечных пользователей для работы со стратегией распространения и без нее, не требуя условной логики. Пример фрагмента кода, иллюстрирующий это:
if tf.config.list_physical_devices('GPU'):
  strategy = tf.distribute.MirroredStrategy()
else:  # use default strategy
  strategy = tf.distribute.get_strategy() 

with strategy.scope():
  # do something interesting
  print(tf.Variable(1.))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>
}

OneDeviceStrategy

tf.distribute.OneDeviceStrategy - это стратегия размещения всех переменных и вычислений на одном указанном устройстве.

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

Эта стратегия отличается от стратегии по умолчанию во многих отношениях. В стратегии по умолчанию логика размещения переменных остается неизменной по сравнению с запуском TensorFlow без какой-либо стратегии распространения. Но при использовании OneDeviceStrategy все переменные, созданные в его области видимости, явно помещаются на указанное устройство. Более того, любые функции, вызываемые через OneDeviceStrategy.run , также будут размещены на указанном устройстве.

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

Подобно стратегии по умолчанию, эту стратегию также можно использовать для тестирования вашего кода перед переключением на другие стратегии, которые фактически распространяются на несколько устройств / машин. Это будет использовать механизм стратегии распространения несколько больше, чем стратегия по умолчанию, но не в полной мере, как при использовании MirroredStrategy или TPUStrategy и т. Д. Если вам нужен код, который ведет себя так, как будто нет стратегии, используйте стратегию по умолчанию.

До сих пор вы видели различные доступные стратегии и способы их реализации. В следующих нескольких разделах показаны различные способы их использования для распределения тренировок.

Использование tf.distribute.Strategy с tf.keras.Model.fit

tf.distribute.Strategy интегрирован в tf.keras который является реализацией tf.keras API . tf.keras - это высокоуровневый API для создания и обучения моделей. Благодаря интеграции в бэкэнд tf.keras мы tf.keras для вас распространение ваших тренировок, написанных в среде обучения model.fit с использованием model.fit .

Вот что вам нужно изменить в своем коде:

  1. Создайте экземпляр соответствующего tf.distribute.Strategy .
  2. Перенесите создание модели Keras, оптимизатора и метрик внутрь strategy.scope .

Мы поддерживаем все типы моделей Keras - последовательные, функциональные и подклассы.

Вот фрагмент кода, чтобы сделать это для очень простой модели Keras с одним плотным слоем:

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).

В этом примере используется MirroredStrategy поэтому вы можете запустить его на машине с несколькими графическими процессорами. strategy.scope() указывает Keras, какую стратегию использовать для распределения обучения. Создание моделей / оптимизаторов / показателей внутри этой области позволяет нам создавать распределенные переменные вместо обычных переменных. После того, как это настроено, вы можете соответствовать своей модели, как обычно. MirroredStrategy заботится о воспроизведении обучения модели на доступных графических процессорах, агрегировании градиентов и многом другом.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)
Epoch 1/2
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
10/10 [==============================] - 3s 2ms/step - loss: 2.3180
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 1.0246
10/10 [==============================] - 1s 2ms/step - loss: 0.5432
0.5432368516921997

Здесьtf.data.Dataset предоставляет данные для обучения иtf.data.Dataset . Вы также можете использовать массивы numpy:

import numpy as np
inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)
Epoch 1/2
10/10 [==============================] - 1s 2ms/step - loss: 0.3866
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.1709
<tensorflow.python.keras.callbacks.History at 0x7f36bfa9fc18>

В обоих случаях (набор данных или numpy) каждый пакет заданного ввода делится поровну между несколькими репликами. Например, при использовании MirroredStrategy с 2 графическими процессорами каждый пакет размером 10 будет разделен между 2 графическими процессорами, каждый из которых получит 5 примеров ввода на каждом этапе. Каждая эпоха будет тренироваться быстрее по мере того, как вы добавляете больше графических процессоров. Как правило, вы хотите увеличивать размер пакета по мере добавления дополнительных ускорителей, чтобы эффективно использовать дополнительную вычислительную мощность. Вам также потребуется перенастроить скорость обучения в зависимости от модели. Вы можете использовать strategy.num_replicas_in_sync чтобы получить количество реплик.

# Compute global batch size using number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
                     mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)

LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

Что сейчас поддерживается?

API обучения ЗеркальныйСтратегия TPUStrategy MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageСтратегия
API Keras Поддерживается Поддерживается Экспериментальная поддержка Экспериментальная поддержка Экспериментальная поддержка

Примеры и учебные пособия

Вот список руководств и примеров, иллюстрирующих приведенную выше интеграцию с Keras:

  1. Учебник по обучению MNIST с помощью MirroredStrategy .
  2. Учебник по обучению MNIST с помощью MultiWorkerMirroredStrategy .
  3. Руководство по обучению MNIST с использованием TPUStrategy .
  4. Учебное пособие для обучения сервера параметров в TensorFlow 2 с ParameterServerStrategy .
  5. Репозиторий TensorFlow Model Garden, содержащий коллекции современных моделей, реализованных с использованием различных стратегий.

Использование tf.distribute.Strategy с настраиваемыми циклами обучения

Как вы видели, использование tf.distribute.Strategy с model.fit требует изменения всего лишь пары строк кода. tf.distribute.Strategy немного больше усилий, вы также можете использовать tf.distribute.Strategy с настраиваемыми циклами обучения.

Если вам нужна большая гибкость и контроль над циклами обучения, чем это возможно с помощью Estimator или Keras, вы можете написать собственные циклы обучения. Например, при использовании GAN вы можете захотеть выполнять разное количество шагов генератора или дискриминатора в каждом раунде. Точно так же высокоуровневые структуры не очень подходят для обучения обучению с подкреплением.

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

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

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

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  optimizer = tf.keras.optimizers.SGD()

Затем создайте входной набор данных и вызовите tf.distribute.Strategy.experimental_distribute_dataset чтобы распределить набор данных на основе стратегии.

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(
    global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

Затем определите один шаг обучения. Используйтеtf.GradientTape для вычисления градиентов и оптимизатора для применения этих градиентов для обновления переменных нашей модели. Чтобы распространить этот шаг обучения, поместите его в функцию train_step и передайте в tf.distrbute.Strategy.run вместе с входными данными набора данных, которые вы получили из созданного ранее dist_dataset :

loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def compute_loss(labels, predictions):
  per_example_loss = loss_object(labels, predictions)
  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    loss = compute_loss(labels, predictions)

  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

@tf.function
def distributed_train_step(dist_inputs):
  per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))
  return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

Еще несколько вещей, на которые следует обратить внимание в приведенном выше коде:

  1. Он использовал tf.nn.compute_average_loss для вычисления потерь. tf.nn.compute_average_loss суммирует потери для каждого примера и делит сумму на global_batch_size. Это важно, потому что позже, после вычисления градиентов для каждой реплики, они агрегируются по репликам путем их суммирования .
  2. Он использовал API tf.distribute.Strategy.reduce для агрегирования результатов, возвращаемых tf.distribute.Strategy.run . tf.distribute.Strategy.run возвращает результаты из каждой локальной реплики в стратегии, и есть несколько способов получить этот результат. Вы можете reduce их, чтобы получить агрегированное значение. Вы также можете выполнить tf.distribute.Strategy.experimental_local_results чтобы получить список значений, содержащихся в результате, по одному на локальную реплику.
  3. Когда apply_gradients вызывается в рамках стратегии распределения, его поведение изменяется. В частности, перед применением градиентов к каждому параллельному экземпляру во время синхронного обучения он выполняет суммирование всех реплик градиентов.

Наконец, как только вы определили шаг обучения, мы можем dist_dataset и запустить обучение в цикле:

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))
tf.Tensor(0.7167484, shape=(), dtype=float32)
tf.Tensor(0.71152556, shape=(), dtype=float32)
tf.Tensor(0.70635474, shape=(), dtype=float32)
tf.Tensor(0.7012355, shape=(), dtype=float32)
tf.Tensor(0.6961673, shape=(), dtype=float32)
tf.Tensor(0.6911497, shape=(), dtype=float32)
tf.Tensor(0.6861821, shape=(), dtype=float32)
tf.Tensor(0.6812641, shape=(), dtype=float32)
tf.Tensor(0.6763952, shape=(), dtype=float32)
tf.Tensor(0.6715747, shape=(), dtype=float32)
tf.Tensor(0.66680235, shape=(), dtype=float32)
tf.Tensor(0.6620776, shape=(), dtype=float32)
tf.Tensor(0.6573999, shape=(), dtype=float32)
tf.Tensor(0.65276873, shape=(), dtype=float32)
tf.Tensor(0.6481838, shape=(), dtype=float32)
tf.Tensor(0.6436445, shape=(), dtype=float32)
tf.Tensor(0.6391503, shape=(), dtype=float32)
tf.Tensor(0.63470083, shape=(), dtype=float32)
tf.Tensor(0.6302957, shape=(), dtype=float32)
tf.Tensor(0.6259341, shape=(), dtype=float32)

В приведенном выше примере вы перебирали dist_dataset чтобы ввести данные для вашего обучения. Мы также предоставляем tf.distribute.Strategy.make_experimental_numpy_dataset для поддержки ввода numpy. Вы можете использовать этот API для создания набора данных перед вызовом tf.distribute.Strategy.experimental_distribute_dataset .

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

iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))
tf.Tensor(0.62161595, shape=(), dtype=float32)
tf.Tensor(0.6173406, shape=(), dtype=float32)
tf.Tensor(0.6131077, shape=(), dtype=float32)
tf.Tensor(0.60891676, shape=(), dtype=float32)
tf.Tensor(0.60476726, shape=(), dtype=float32)
tf.Tensor(0.60065883, shape=(), dtype=float32)
tf.Tensor(0.59659094, shape=(), dtype=float32)
tf.Tensor(0.5925633, shape=(), dtype=float32)
tf.Tensor(0.5885753, shape=(), dtype=float32)
tf.Tensor(0.5846267, shape=(), dtype=float32)

Это охватывает простейший случай использования tf.distribute.Strategy API для распределения пользовательских циклов обучения.

Что сейчас поддерживается?

API обучения ЗеркальныйСтратегия TPUStrategy MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageСтратегия
Индивидуальный цикл обучения Поддерживается Поддерживается Экспериментальная поддержка Экспериментальная поддержка Экспериментальная поддержка

Примеры и учебные пособия

Вот несколько примеров использования стратегии распространения с настраиваемыми циклами обучения:

  1. Учебник по обучению MNIST с помощью MirroredStrategy .
  2. Руководство по обучению MNIST с использованием TPUStrategy .
  3. Репозиторий TensorFlow Model Garden, содержащий коллекции современных моделей, реализованных с использованием различных стратегий.

Другие темы

В этом разделе рассматриваются некоторые темы, относящиеся к нескольким вариантам использования.

Настройка переменной окружения TF_CONFIG

Для обучения с несколькими рабочими, как упоминалось ранее, вам необходимо установить переменную среды TF_CONFIG для каждого двоичного TF_CONFIG запущенного в вашем кластере. TF_CONFIG среды TF_CONFIG - это строка JSON, которая указывает, какие задачи составляют кластер, их адреса и роль каждой задачи в кластере. Репозиторий тензорного потока / экосистемы предоставляет шаблон Kubernetes, в котором задает TF_CONFIG для ваших обучающих задач.

TF_CONFIG состоит из двух компонентов: кластера и задачи. cluster предоставляет информацию об обучающем кластере, который представляет собой диктатор, состоящий из разных типов заданий, таких как worker. В обучении с несколькими работниками обычно есть один работник, который берет на себя немного больше ответственности, например, сохраняет контрольную точку и написание сводного файла для TensorBoard в дополнение к тому, что делает обычный работник. Такой рабочий называется «главным» рабочим, и обычно рабочий с индексом 0 назначается главным рабочим (на самом деле так реализуется tf.distribute.Strategy). задача, с другой стороны, предоставляет информацию о текущей задаче. Первый кластер компонентов одинаков для всех рабочих процессов, а задача второго компонента отличается для каждого рабочего и определяет тип и индекс этого рабочего.

Один из TF_CONFIG :

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
        "ps": ["host4:port", "host5:port"]
    },
   "task": {"type": "worker", "index": 1}
})

Этот TF_CONFIG указывает, что в TF_CONFIG есть три рабочих и две задачи ps, а также их хосты и порты. Часть «задача» указывает, что роль текущей задачи в кластере - worker 1 (второй worker). Допустимые роли в кластере: «руководитель», «работник», «пс» и «оценщик». tf.distribute.experimental.ParameterServerStrategy "ps" быть не должно, кроме случаев использования tf.distribute.experimental.ParameterServerStrategy .

Что дальше?

tf.distribute.Strategy находится в стадии активной разработки. Попробуйте и поделитесь своими отзывами с помощью GitHub .