Статьи TensorFlow | Обзор распределенного обучения TensorFlow 2.x

TensorFlow

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

Распределенный режим обучения

В соответствии с различными проблемами, с которыми сталкивается модель при обучении одного узла, ее распределенное обучение можно разделить на параллелизм модели (Model Parallelism) и параллелизм данных (Data Parallelism) в двух режимах.

модельный параллелизм

модельный параллелизм是指将模型切分为多个部分并将各个部分放置到不同的节点进行训练的分布式模式. Потому что, когда модель очень сложная и много параметров, один узел не может загрузить и обучить всю модель из-за ограничений памяти, поэтому необходимо разрезать модель на более мелкие части и запускать каждую часть на разных узлах для завершения обучения. .

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

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

Поэтому, если модель не очень велика или ее части могут выполняться параллельно, обычно не рекомендуется использовать режим распределенного обучения параллелизма модели. существуетTensorFlowМодельный параллелизм также известен как репликация внутри графа (In-graph Replication).

параллелизм данных

параллелизм данных是指将输入数据进行切分并使用不同节点的模型副本进行训练的分布式模式. Потому что, когда объем данных велик, выполнение одного узла занимает много времени.1раунд обучения, что невыносимо для десятков или десятков раундов обучения, поэтому данные могут быть разделены и распределены по нескольким узлам, каждый узел обрабатывает свою небольшую часть данных, тем самым ускоряя общую скорость обучения. существуетTensorFlowПараллелизм данных также известен как репликация между графами (Between-graph Replication).

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

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

  2. Асинхронное обучение означает, что каждый узел1КусокbatchПосле обучения вам не нужно ждать, пока другие узлы обновят значение переменных модели напрямую. Скорость обучения асинхронного обучения выше, но у него есть серьезная проблема отказа градиента (Stale Gradients), потому что узел, который завершит обучение позже, обнаружит, что значение переменной модели было обновлено другими узлами, и градиент, рассчитанный этим узлом, истечет. Из-за проблемы отказа градиента модель может в конечном итоге попасть в неоптимальное решение (Sub-Optimal Training).

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

Распределенная архитектура обучения

Представленная здесь архитектура распределенного обучения в основном относится к数据并行Распределенная архитектура в режиме. В настоящее время в основном используются два типа распределенных обучающих архитектур, а именно:Parameter Serverархитектура иAllReduceАрхитектура.

Архитектура PS

Parameter Server(PS)Схема относится к серверу с параметрами (ps) распределенная кластерная архитектура. В этой архитектуре некоторые узлы в кластере действуют как серверы параметров для хранения и обновления значений переменных модели, в то время как другие узлы используются для вычисления градиентов переменных модели и взаимодействия с сервером параметров.

Parameter ServerАрхитектура может быть либо синхронным режимом обучения, либо асинхронным режимом обучения.

В режиме синхронного обученияpsузел будет ждать всехworkerодин из узловbatchобучение заканчивается, затем объединяем этиworkerГрадиенты переменных модели для узлов и обновленияpsЗначение переменной модели в , и, наконец, транслировать обновленное значение всемworkerузлы для обновления своих переменных модели.

В асинхронном режиме обучения всеworkerУзлы обновляются индивидуальноpsПеременная в узле, а затем получить последнее значение переменной, чтобы продолжить обучение.

Однако, будь то синхронный или асинхронный режим обучения, в этой архитектуреpsПропускная способность узлов может стать узким местом для обучения модели, потому что когдаworkerКогда число узлов увеличивается,psДавление пропускной способности на узле будет увеличиваться экспоненциально. Хотя можно настроитьpsКоличество узлов для уменьшения нагрузки на полосу пропускания, ноpsСтоимость связи между узлами.

Архитектура AllReduce

AllReduceАрхитектура относится к архитектуре распределенного кластера без серверов параметров. В этой архитектуре все узлы в кластере действуют какworkerдля выполнения вычислительных операций архитектура будетbatchИспользовать после тренировкиAllReduceалгоритм во всемworkerСинхронное обновление переменных модели между узлами.

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

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

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

Стратегия распределенного обучения TF

TensorFlowизtf.distributeМодуль содержит серию распределенных стратегий обучения, основанных на数据并行режим реализован. Некоторые стратегии до сих порexperimentalПод модулем имеется в виду, что они являются экспериментальными стратегиями и могут измениться в будущем, но в основном их функции относительно зрелые и законченные, и вы можете попробовать их использовать. учитываяTPUStrategyа такжеOneDeviceStrategyОсобые сценарии использования , здесь подробно описываться не будут.

Mirrored

MirroredStrategyЭто одномашинная синхронная распределенная стратегия обучения. Он поддерживает несколькоGPUРаспределенное обучение междуGPUсоздает копию модели на модели с каждой переменной в модели (Variables) будет зеркально отображен и помещен в соответствующийGPU, эти переменные называются зеркальными переменными (MirroredVariable).

MirroredStrategyстратегия принятаAllReduceАлгоритм делает все зеркальные переменные в каждомGPUпродолжайте синхронизировать обновления междуAllReduceПо умолчанию используется алгоритм NVIDIA.NcclAllReduce, также черезcross_device_opsИзмените параметр на другойAllReduceалгоритмы, такие какHierarchicalCopyAllReduce.

MirroredStrategyПолитика автоматически использует все доступныеTensorFlowнашелGPUпроводить распределенное обучение, если вы хотите использовать только частьGPUтогда вы можете пройтиdevicesпараметры указать.MirroredStrategyКод создания экземпляра выглядит следующим образом:

mirrored_strategy = tf.distribute.MirroredStrategy(
    devices=["/gpu:0", "/gpu:1"],
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce(),
)

еслиTensorFlowне найденGPUПо умолчанию он вырождается в использованиеCPUтренировать.MirroredStrategyТипичный сценарий использования — одна машина с несколькимиGPU.

CentralStorage

CentralStorageStrategyЭто также стратегия синхронного распределенного обучения на одной машине. но сMirroredStrategyСтратегия отличается тем, что сохраняет все переменные модели вCPUв памяти, а не путем зеркального копирования в каждомGPU, все вычислительные операции будут выполняться на каждомGPUвыполняется так же, как указано выше.

Если машина имеет только одинGPU, то все переменные и вычисления будут размещены в этомGPUначальство. справаCPUПеред обновлением переменных вGPUПеременные градиенты на репликах агрегируются, а затем применяются кCPUПеременные обновляются.CentralStorageStrategyКод создания экземпляра выглядит следующим образом:

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()

CentralStorageStrategyстратегия вCPUиGPUСтоимость связи намного ниже, чемGPUиGPUЭто более применимо, когда стоимость связи между ними.

MultiWorkerMirrored

MultiWorkerMirroredStrategyстратегия иMirroredStrategyСтратегии очень похожи и могут быть поняты какMirroredStrategyстратегический多机的Синхронизированная версия распределенного обучения, которая также создает копии всех переменных на каждой машине.

MultiWorkerMirroredStrategyКаждый узел, работающий в политике, называетсяworker,ДолженworkerУзел может содержать ноль или болееGPU. несколькоworkerиспользовать между узламиAllReduceалгоритм для синхронного обновления переменных модели,TensorFlowЭта операция называетсяCollectiveOps.CollectiveOpsБудет вTensorFlowКогда модель запускается, она автоматически выбирает подходящую на основе оборудования, топологии сети и размера тензора.AllReduceАлгоритмы для связи по сети для завершения обновления переменных.

MultiWorkerMirroredStrategyСтратегия также реализует некоторые дополнительные оптимизации производительности, такие как объединение нескольких небольших тензоров.AllReductionsПреобразовано в большой тензорAllReductions,УменьшитьworkerВремя связи между узлами.

MultiWorkerMirroredStrategyНа данный момент есть две стратегии на выборCollectiveOps. дляCollectiveCommunication.RING, который используетgRPCВ качестве коммуникационного слоя, реализованного на основе环изAllReduceработать. ДругойCollectiveCommunication.NCCL, который использует NvidiaNCCLбиблиотека для реализацииAllReduceработать. При фактическом использовании вы можете выбрать подходящий вариант в зависимости от вашей операционной среды.CollectiveOpsили используйтеCollectiveCommunication.AUTOсдаватьTensorFlowВыбирайте во время выполнения.MultiWorkerMirroredStrategyКод создания экземпляра выглядит следующим образом:

multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
    tf.distribute.experimental.CollectiveCommunication.RING)

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

TensorFlowСтандартный способ определить информацию о конфигурации кластера — использоватьTF_CONFIGОн реализуется переменными среды, которые определяют информацию о конфигурации всех узлов в кластере, включая всеworkerСетевой адрес узла, в настоящее времяworkerиндекс узла (index) и текущийworkerРоль узла (type).

я упалworkerузел не содержитGPU, стратегия сводится к использованиюCPUв несколькихworkerРаспределенное обучение среди узлов. если кластерworkerЕсли число узлов равно одному, стратегия вырождается вMirroredStrategyСтратегия.

ParameterServer

ParameterServerStrategyЭто многомашинная асинхронная распределенная стратегия обучения. Так что это тоже нужно уточнять заранееTF_CONFIGинформация о переменной среды иMultiWorkerMirroredStrategyРазница в стратегии заключается в том, что не все узлы в кластереworker, некоторые узлы будут обозначены какpsИспользуется для хранения переменной информации. Каждая переменная модели хранится вpsНа узле все вычислительные операции будут выполняться на всехworkerУзлы выполняются таким же образом.ParameterServerStrategyКод создания экземпляра выглядит следующим образом:

ps_strategy = tf.distribute.experimental.ParameterServerStrategy()

распределенный кластер TF

определение кластера

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

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

типичныйTF_CONFIGЗначения переменных окружения следующие:

{
  "cluster": {
    "chief": ["host1:port"],
    "worker": ["host2:port", "host3:port"],
    "ps": ["host4:port"],
    "evaluator": ["host5:port"]
  },
  "task": {
    "type": "worker",
    "index": 0
  }
}

этоjsonнить.clusterИнформация указывает на то, что кластер1Кусокchiefузел,2Кусокworkerузел,1Кусокpsузел и1Кусокevaluatorузел,taskИнформация указывает на то, что текущий узел является первым1Кусокworkerузел.

chiefРоль узлов иworkerNode примерно такой же, но он выполняет дополнительную работу, например, сохраняет файлы контрольных точек (checkpoints) и дляTensorboardфайл журнала и т. д., если не указаноcheifузел, по умолчанию будетworkerпервый узел в списке какchiefузел;workerУзлы используются для выполнения обучающих операций;psУзлы используются для хранения переменных, только при использованииParameterServerStrategyЕго нужно указать только при обучении стратегии;evaluatorИспользуется для выполнения операций перекрестной проверки, обычно также используетсяParameterServerStrategyуказана политика.

Обратите внимание, что все узлыTF_CONFIGв переменной окруженияclusterИнформация та же, разница в том, чтоtaskчасть, и все роли(task type) изindexдолжно быть из0начать, потому чтоTensorFlowбудет основываться наindexотclusterПрочитайте информацию об узле из списка под соответствующей ролью.

TF_CONFIGПеременная среды может быть записана в переменную среды системы, но предполагается, что только один экземпляр узла кластера будет запущен на физическом узле одновременно, но на практике мы можем запустить несколько экземпляров узла на одном и том же физическом узле. узел, использующий разные порты.TF_CONFIGзначение потеряет гибкость. Поэтому в большинстве случаев мы будемpythonчерез программуos.environ["TF_CONFIG"]указать информацию о кластере для создания по запросу,TensorFlowСреда выполнения автоматически анализирует информацию и запускает задачу обучения.

описано выше4Среди распределенных стратегий обучения в Китае для многомашинного обучения можно использовать только кластер.MultiWorkerMirroredStrategyиParameterServerStrategyСтратегия.

Распределенное обучение кластера TF

TensorFlowБазовый процесс кластерного распределенного обучения выглядит следующим образом:

  1. Разные узлы используют разныеTF_CONFIGПеременная среды запускает задачу обучения для этого узла.TensorFlowбудет автоматически проанализированTF_CONFIGзначение переменной среды, затем используйтеip:portзапустить текущий узелgRPCобслуживание и прослушивание других узловgRPCСлужить.

  2. когда все узлыgRPCПосле того, как услуги готовы, каждыйworkerУзлы начинают обучение со своими собственными данными, и начинается глобальное обучение.

  3. КаждыйbatchПо окончании обучения,TensorFlowВсе будет обновляться в соответствии с различными распределенными стратегиями.workerПеременная узла, после завершения обновления перейти к следующемуbatchобучение.

  4. После выполнения всех операций обучения и оценки обучение заканчивается, иgRPCПосле этого служба отключается.

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

Итак, нам нужно найти автоматическую сборкуTF_CONFIGМетод переменных окружения, некоторые распределенные обучающие фреймворки могут решить наши проблемы. Например360РелизXLearningплатформа, используяYarnДля управления и распределения ресурсов кластера каждый раз, когда запускается задача обучения, она автоматически выделяет указанное количество ресурсов и определяет сборку.TF_CONFIGИнформация, требуемая переменными среды, нам нужно только проанализировать информацию, чтобы начать обучение, что может значительно снизить рабочую нагрузку на ранней стадии обучения, чтобы мы могли сосредоточиться на конкретном процессе обучения.

XLearningэто всего лишь одна из многих распределенных обучающих сред, и некоторые из них похожи наGoogleизKubeflowПлатформы и тому подобное также могут реализовывать аналогичные функции. и их роль не ограничиваетсяTF_CONFIGПостроение переменных среды также может быть выполненоCPUТакие операции, как выделение ядра, выделение памяти и разделение данных. Использование этих фреймворков позволяет нам использоватьTensorFlowБолее удобно при распределенном обучении.

Реализация распределенного обучения TF

TensorFlowПоддерживает использование более высокого порядкаAPIРаспределенное обучение может быть реализовано в виде распределенного обучения, а распределенное обучение также может быть реализовано с использованием пользовательского процесса обучения.

Keras API

KerasдаTensorFlowПостроить более высокий порядок моделиAPI, который был объединен сTensorFlowСтратегия распределенного обучения легко интегрируется, и это также официально рекомендуемый предпочтительный способ достижения распределенного обучения.

описано выше4В распределенной стратегии использованиеKeras APIСтратегии распределенного обучения:MirroredStrategy,CentralStorageStrategyа такжеMultiWorkerMirroredStrategy, иParameterServerStrategyОжидается, что политики будут поддерживаться в последующих выпусках. Итак, в настоящее время используетсяKerasПри выполнении распределенного обучения на нескольких машинах можно использовать толькоMultiWorkerMirroredStrategyраспределенная стратегия.

использоватьKerasОбучающий код построенной локальной модели можно легко преобразовать в соответствующий распределенный обучающий код. Просто нужно построить и скомпилировать модель (compile) операции размещаются в экземпляре распределенной политикиscopeОн может поддерживать все типыKerasмодели, в том числеSequential,Functionalа такжеSubclassed. использоватьKerasПример фрагмента кода для распределенного обучения показан ниже:

import tensorflow as tf
import numpy as np

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1, ))])
    model.compile(loss='mse', optimizer='sgd')

inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)

вstrategyможет быть поддержанKerasЛюбой экземпляр политики для распределенного обучения.strategy.scope()Указывает блок кода, который будет выполняться распределенно, обратите вниманиеfitМетоды не могут быть включены.fitОперации могут быть согласованы с местным обучением без дополнительных модификаций.

использоватьKerasПри выполнении распределенного обучения распределенная стратегия обучения и оценки является последовательной, то есть операции оценки модели также болееworkerУзлы совместно участвуют в выполнении. заKerasвходные данные в модель,TensorFlowБудут использоваться различные стратегии для автоматического разделения данных в соответствии с источником данных (текстовый файл или числовое значение), чтобы каждыйworkerЧасть общих данных обучения ноды, конечно, тоже можно пропустить черезtf.data.OptionsКонтролировать, чтобы входные данные не были разделены.

Estimator API

EstimatorдаTensorFlowДругой более высокий порядок построения моделейAPI, который изначально поддерживает асинхронныйParameter ServerРаспределенная стратегия обучения. иKerasаналогично, это такжеTensorFlowРаспределенные стратегии обучения интегрированы, но с ограниченной поддержкой. использоватьEstimatorПостроенный код можно обучать как локально, так и распространять без дополнительных модификаций. При распределенном обучении используется распределенная стратегия по умолчанию.ParameterServerStrategy.

иKerasиспользуется вstrategy.scopeРаспределенное обучение отличается,Estimatorпоставивstrategyэкземпляр переданRunConfigиспользовать различные распределенные стратегии обучения.EstimatorМодуль содержит некоторые предопределенные модели, которые можно использовать напрямую, такие как модели линейной регрессии и т. д. Если вы хотите настроить модель, вам необходимо сначала определить конструктор модели.model_fn.

Здесь предопределеноEstimatorМодельLinearRegressorНапример, чтобы представить его с помощьюMirroredStrategyМетод реализации стратегии распределенного обучения, код примера выглядит следующим образом:

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
config = tf.estimator.RunConfig(
    train_distribute=strategy,
    eval_distribute=strategy,
)
regressor = tf.estimator.LinearRegressor(
    feature_columns=[tf.feature_column.numeric_column('feats')],
    optimizer='SGD',
    config=config,
)

def input_fn():
    dataset = tf.data.Dataset.from_tensors(({"feats": [1.]}, [1.]))
    return dataset.repeat(1000).batch(10)
tf.estimator.train_and_evaluate(
    regressor,
    train_spec=tf.estimator.TrainSpec(input_fn=input_fn),
    eval_spec=tf.estimator.EvalSpec(input_fn=input_fn),
)

Примечание сKerasразные,Estimatorв состоянии пройтиRunConfigметодtrain_distributeиeval_distributeПараметры определяют распределенную стратегию обучения и оценки отдельно, что является более гибким.

и используяEstimatorПри выполнении распределенного обучения входные данные модели не будут автоматически разделяться, и она полностью независима и управляема.Можно настроить метод обработки входных данных.Например, мы можем использоватьdataset.shardметод разделения или для данных, считанных из файла, вы можете заранее разделить файл и назначить его различнымworkerузел.

Следует отметить, что если всеworkerВсе узлы обучаются на одних и тех же данных, рекомендуются разные семена (seed) для входных данныхdatasetпровестиshuffle, чтобы избежать градиентного спуска всех узлов в одном направлении, что приводит к увеличению времени сходимости и ухудшению конечных результатов обучения.

В настоящее времяTensorFlow 2.xРаспределенная стратегия обучения дляEstimatorПоддержка относительно ограничена, базовое обучение и оценка могут выполняться в обычном режиме, но некоторые расширенные функции, такие какscaffoldи т. д. не поддерживаются, а также могут содержать некоторые потенциальныеbug. В краткосрочной перспективе официальная команда не планирует тратить слишком много усилий на улучшениеEstimatorраспределенное обучение, они больше сосредоточатся наKerasА индивидуальный распределенный метод обучения обеспечивает техническую поддержку, поэтому здесь не рекомендуетсяEstimatorдля распределенного обучения.

Индивидуальное распределенное обучение

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

можно использоватьtf.distribute.StrategyРяд основных методов класса для поддержки пользовательского распределенного процесса обучения. сначала сKerasРаспределенное обучение похоже наstrategy.scope()Создайте модель и оптимизатор под (optimizers); затем создайте входные данные и используйтеtf.distribute.Strategy.experimental_distribute_datasetметод для размещения данных; затем используйтеtf.GradientTape()Определяет поток обучения; наконец, перебирает входные данные и начинает обучение. Пример кода выглядит следующим образом:

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
global_batch_size = 128
with strategy.scope():
    model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1, ))])
    optimizer = tf.keras.optimizers.SGD()

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

@tf.function
def train_step(dist_inputs):
    def step_fn(inputs):
        features, labels = inputs

        with tf.GradientTape() as tape:
            logits = model(features, training=True)
            cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
                logits=logits,
                labels=labels,
            )
            loss = tf.reduce_sum(cross_entropy) * (1.0 / global_batch_size)

        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
        return cross_entropy

    per_example_losses = strategy.run(step_fn, args=(dist_inputs, ))
    mean_loss = strategy.reduce(
        tf.distribute.ReduceOp.MEAN,
        per_example_losses,
        axis=0,
    )
    return mean_loss

with strategy.scope():
    for inputs in dist_dataset:
        print(train_step(inputs))

Обратите внимание, что код используетstrategy.runвыполнитьbatchобучение, использованиеstrategy.reduceдля агрегирования результатов обучения каждого узла. когда будетapply_gradientsметод помещается вstrategy.scope()Когда в следующий раз его поведение при обновлении немного изменится по сравнению с локальным обучением, он сначала агрегирует всеworkerГрадиенты узлов суммируются, а затем это значение используется для обновления переменных модели каждого узла.

использованная литература

  1. Официальная документация по распределенному обучению TensorFlow
  2. Принцип и практика распределенного обучения TensorFlow
  3. Начало работы с распределенным TensorFlow
  4. Прошлое и настоящее алгоритма AllReduce
  5. Адрес проекта XLearning
Категории