Есть вопрос? Присоединяйтесь к сообществу на форуме TensorFlow. Посетите форум.

Распределенное обучение с 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.

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

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

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

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

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

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

mirrored_strategy = tf.distribute.MirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
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:1,/job:localhost/replica:0/task:0/device:GPU:0
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
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())
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
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:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.AUTO

MultiWorkerMirroredStrategy имеет две реализации для обмена данными между устройствами. CommunicationImplementation.RING основан на RPC и поддерживает как процессоры, так и графические процессоры. CommunicationImplementation.NCCL использует NCCL и обеспечивает высочайшую производительность на графических процессорах, но не поддерживает процессоры. 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: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)

В TensorFlow 1 ParameterServerStrategy доступен только с tf.compat.v1.distribute.experimental.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 the Default Strategy
  strategy = tf.distribute.get_strategy() 

with strategy.scope():
  # Do something interesting
  print(tf.Variable(1.))
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
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 вы можете легко распространять свои тренировки, написанные в среде обучения Model.fit с помощью Model.fit .

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

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

Стратегии распространения TensorFlow поддерживают все типы моделей 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')
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
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: 0.3648
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: 0.1613
10/10 [==============================] - 1s 2ms/step - loss: 0.1001
0.1001412644982338

Здесь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.0713
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.0315
<tensorflow.python.keras.callbacks.History at 0x7f18ac0bdcd0>

В обоих случаях - с Dataset или 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.5716536, shape=(), dtype=float32)
tf.Tensor(0.56787133, shape=(), dtype=float32)
tf.Tensor(0.564126, shape=(), dtype=float32)
tf.Tensor(0.5604175, shape=(), dtype=float32)
tf.Tensor(0.5567452, shape=(), dtype=float32)
tf.Tensor(0.5531087, shape=(), dtype=float32)
tf.Tensor(0.54950774, shape=(), dtype=float32)
tf.Tensor(0.5459418, shape=(), dtype=float32)
tf.Tensor(0.5424106, shape=(), dtype=float32)
tf.Tensor(0.5389137, shape=(), dtype=float32)
tf.Tensor(0.53545076, shape=(), dtype=float32)
tf.Tensor(0.5320213, shape=(), dtype=float32)
tf.Tensor(0.528625, shape=(), dtype=float32)
tf.Tensor(0.5252616, shape=(), dtype=float32)
tf.Tensor(0.52193063, shape=(), dtype=float32)
tf.Tensor(0.5186317, shape=(), dtype=float32)
tf.Tensor(0.5153646, shape=(), dtype=float32)
tf.Tensor(0.51212865, shape=(), dtype=float32)
tf.Tensor(0.5089238, shape=(), dtype=float32)
tf.Tensor(0.5057497, 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.5026059, shape=(), dtype=float32)
tf.Tensor(0.49949214, shape=(), dtype=float32)
tf.Tensor(0.496408, shape=(), dtype=float32)
tf.Tensor(0.4933532, shape=(), dtype=float32)
tf.Tensor(0.49032742, shape=(), dtype=float32)
tf.Tensor(0.48733026, shape=(), dtype=float32)
tf.Tensor(0.48436147, shape=(), dtype=float32)
tf.Tensor(0.4814206, shape=(), dtype=float32)
tf.Tensor(0.4785076, shape=(), dtype=float32)
tf.Tensor(0.475622, 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, которая определяет, какие задачи составляют кластер, их адреса и роль каждой задачи в кластере. tensorflow/ecosystem предоставляет шаблон Kubernetes, который настраивает TF_CONFIG для ваших учебных задач.

TF_CONFIG из двух компонентов: кластера и задачи.

  • Кластер предоставляет информацию об обучающем кластере, который состоит из разных типов заданий, таких как 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 указывает, что в "cluster" есть три рабочих и две задачи "ps" "cluster" а также их хосты и порты. Часть "task" указывает, что роль текущей задачи в "cluster" работник 1 (второй работник). Допустимые роли в кластере: "chief" , "worker" , "ps" и "evaluator" . tf.distribute.experimental.ParameterServerStrategy "ps" быть не должно, кроме случаев использования tf.distribute.experimental.ParameterServerStrategy .

Что дальше?

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