Ta strona została przetłumaczona przez Cloud Translation API.
Switch to English

Rozproszone szkolenie z TensorFlow

Zobacz na TensorFlow.org Uruchom w Google Colab Wyświetl źródło na GitHub Pobierz notatnik

Przegląd

tf.distribute.Strategy to TensorFlow API do dystrybucji szkoleń na wiele GPU, wiele maszyn lub TPU. Korzystając z tego interfejsu API, możesz rozpowszechniać istniejące modele i kod szkoleniowy przy minimalnych zmianach kodu.

tf.distribute.Strategy została zaprojektowana z myślą o następujących kluczowych celach:

  • Łatwy w użyciu i obsługuje wiele segmentów użytkowników, w tym badaczy, inżynierów ML itp.
  • Zapewniają dobrą wydajność po wyjęciu z pudełka.
  • Łatwe przełączanie między strategiami.

tf.distribute.Strategy może być używany z interfejsem API wysokiego poziomu, takim jak Keras , a także może być używany do dystrybucji niestandardowych pętli szkoleniowych (i ogólnie wszelkich obliczeń wykorzystujących TensorFlow).

W TensorFlow 2.x możesz uruchamiać programy z zapałem lub na wykresie za pomocą tf.function . tf.distribute.Strategy zamierza obsługiwać oba te tryby wykonywania, ale najlepiej działa z tf.function . Tryb Eager jest zalecany tylko do debugowania i nie jest obsługiwany przez TPUStrategy . Chociaż szkolenie jest głównym tematem tego przewodnika, ten interfejs API może być również używany do rozpowszechniania oceny i przewidywania na różnych platformach.

Możesz używać tf.distribute.Strategy z bardzo niewielkimi zmianami w kodzie, ponieważ zmieniliśmy podstawowe komponenty TensorFlow, aby były świadome strategii. Obejmuje to zmienne, warstwy, modele, optymalizatory, metryki, podsumowania i punkty kontrolne.

W tym przewodniku wyjaśniamy różne rodzaje strategii i wyjaśniamy, jak można ich używać w różnych sytuacjach. Aby dowiedzieć się, jak debugować problemy z wydajnością, zobacz przewodnik dotyczący optymalizacji wydajności GPU TensorFlow .

# Import TensorFlow
import tensorflow as tf

Rodzaje strategii

tf.distribute.Strategy zamierza objąć szereg przypadków użycia na różnych osiach. Niektóre z tych kombinacji są obecnie obsługiwane, a inne zostaną dodane w przyszłości. Niektóre z tych osi to:

  • Trenowanie synchroniczne i asynchroniczne: są to dwa typowe sposoby dystrybucji szkolenia z równoległością danych. W przypadku treningu synchronicznego wszyscy pracownicy szkolą się na różnych fragmentach danych wejściowych w synchronizacji i agregując gradienty na każdym kroku. W szkoleniu asynchronicznym wszyscy pracownicy niezależnie uczą się danych wejściowych i asynchronicznie aktualizują zmienne. Zwykle szkolenie w zakresie synchronizacji jest obsługiwane przez redukcję wszystkich elementów i asynchronizację poprzez architekturę serwera parametrów.
  • Platforma sprzętowa: możesz chcieć skalować swoje szkolenie na wiele procesorów graficznych na jednej maszynie lub na wiele maszyn w sieci (z co najmniej 0 procesorów graficznych w każdym) lub na Cloud TPU.

Aby wspierać te przypadki użycia, dostępnych jest sześć strategii. Następna sekcja wyjaśnia, które z nich są obsługiwane w jakich scenariuszach w FT. Oto krótki przegląd:

Szkoleniowy interfejs API MirroredStrategy Strategia TPUS MultiWorkerMirroredStrategy CentralStorageStrategy ParameterServerStrategy
Keras API Utrzymany Utrzymany Wsparcie eksperymentalne Wsparcie eksperymentalne Obsługiwane planowane stanowisko 2.3
Niestandardowa pętla treningowa Utrzymany Utrzymany Wsparcie eksperymentalne Wsparcie eksperymentalne Obsługiwane planowane stanowisko 2.3
Estimator API Ograniczone wsparcie Niewspierany Ograniczone wsparcie Ograniczone wsparcie Ograniczone wsparcie

MirroredStrategy

tf.distribute.MirroredStrategy obsługuje synchroniczne rozproszone szkolenie na wielu procesorach graficznych na jednym komputerze. Tworzy jedną replikę na każde urządzenie GPU. Każda zmienna w modelu jest odzwierciedlana we wszystkich replikach. Zmienne te razem tworzą pojedynczą zmienną koncepcyjną o nazwie MirroredVariable . Te zmienne są ze sobą zsynchronizowane, stosując identyczne aktualizacje.

Do przekazywania aktualizacji zmiennych między urządzeniami używane są wydajne algorytmy all-RED. Zmniejsz wszystkie agregaty tensorów na wszystkich urządzeniach, dodając je i udostępniaj na każdym urządzeniu. Jest to połączony algorytm, który jest bardzo wydajny i może znacznie zmniejszyć obciążenie związane z synchronizacją. Dostępnych jest wiele algorytmów i implementacji typu all-RED, w zależności od typu komunikacji dostępnej między urządzeniami. Domyślnie używa NVIDIA NCCL jako implementacji all-RED. Możesz wybrać jedną z kilku innych opcji lub napisać własną.

Oto najprostszy sposób tworzenia MirroredStrategy :

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

Spowoduje to utworzenie instancji MirroredStrategy która będzie używać wszystkich procesorów graficznych widocznych dla TensorFlow i używać NCCL jako komunikacji między urządzeniami.

Jeśli chcesz używać tylko niektórych procesorów graficznych na swoim komputerze, możesz to zrobić w następujący sposób:

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
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')

Jeśli chcesz przesłonić komunikację między urządzeniami, możesz to zrobić za pomocą argumentu cross_device_ops , dostarczając instancję tf.distribute.CrossDeviceOps . Obecnie tf.distribute.HierarchicalCopyAllReduce i tf.distribute.ReductionToOneDevice to dwie opcje inne niż tf.distribute.NcclAllReduce które jest wartością domyślną.

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',)

Strategia TPUS

tf.distribute.TPUStrategy umożliwia prowadzenie szkolenia TensorFlow na jednostkach przetwarzania tensorowego (TPU). TPU to wyspecjalizowane układy ASIC firmy Google zaprojektowane w celu radykalnego przyspieszenia obciążeń systemów uczących się. Są dostępne w Google Colab, TensorFlow Research Cloud i Cloud TPU .

Pod względem architektury rozproszonego szkolenia TPUStrategy jest tą samą MirroredStrategy - implementuje synchroniczne rozproszone szkolenie. Jednostki TPU zapewniają własną implementację wydajnych operacji all-REDUCH i innych zbiorczych operacji na wielu rdzeniach TPU, które są wykorzystywane w TPUStrategy .

Oto jak możesz utworzyć instancję 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)

Instancja TPUClusterResolver pomaga zlokalizować jednostki TPU. W Colab nie musisz podawać żadnych argumentów.

Jeśli chcesz użyć tego dla Cloud TPU:

  • Musisz określić nazwę swojego zasobu TPU w argumencie tpu .
  • Musisz jawnie zainicjować system tpu na początku programu. Jest to wymagane, zanim jednostki TPU będą mogły być używane do obliczeń. Inicjalizacja systemu tpu powoduje również wyczyszczenie pamięci TPU, dlatego ważne jest, aby najpierw wykonać ten krok, aby uniknąć utraty stanu.

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy jest bardzo podobny do MirroredStrategy . Wdraża synchroniczne, rozproszone szkolenie dla wielu pracowników, z których każdy może mieć wiele procesorów graficznych. Podobnie jak tf.distribute.MirroredStrategy , tworzy kopie wszystkich zmiennych w modelu na każdym urządzeniu dla wszystkich pracowników.

Oto najprostszy sposób tworzenia 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 ma dwie implementacje do komunikacji między urządzeniami. CommunicationImplementation.RING jest oparty na RPC i obsługuje zarówno CPU, jak i GPU. CommunicationImplementation.NCCL wykorzystuje NCCL Nvidii i zapewnia najnowocześniejszą wydajność na GPU, ale nie obsługuje procesora. CollectiveCommunication.AUTO odkłada wybór na Tensorflow. Możesz je określić w następujący sposób:

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

Jedną z kluczowych różnic między szkoleniem dla wielu pracowników w porównaniu ze szkoleniem z wieloma GPU jest konfiguracja dla wielu pracowników. TF_CONFIG środowiskowa TF_CONFIG jest standardowym sposobem w TensorFlow do określania konfiguracji klastra dla każdego procesu roboczego, który jest częścią klastra. Dowiedz się więcej o konfigurowaniu TF_CONFIG .

ParameterServerStrategy

Trenowanie serwera parametrów jest popularną metodą równoległą do danych w celu skalowania szkolenia modelu na wielu komputerach. Klaster szkoleniowy serwera parametrów składa się z procesów roboczych i serwerów parametrów. Zmienne są tworzone na serwerach parametrów i są odczytywane i aktualizowane przez pracowników na każdym kroku. Szczegółowe informacje można znaleźć w samouczku szkoleniowym dotyczącym serwera parametrów .

Trening parametrów serwera TensorFlow 2 wykorzystuje architekturę opartą na centralnym koordynatorze za pośrednictwem klasy tf.distribute.experimental.coordinator.ClusterCoordinator .

W tej implementacji, że worker i parameter server zadania uruchamiane tf.distribute.Server jakoby słuchać zadań z koordynatorem. Koordynator tworzy zasoby, wysyła zadania szkoleniowe, pisze punkty kontrolne i radzi sobie z awariami zadań.

W programowaniu działającym na koordynatorze użyjesz obiektu ParameterServerStrategy do zdefiniowania kroku szkoleniowego i użyjesz ClusterCoordinator do wysłania kroków szkoleniowych do pracowników zdalnych. Oto najprostszy sposób ich tworzenia:

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

Zauważ, że będziesz musiał skonfigurować zmienną środowiskową TF_CONFIG, jeśli używasz TFConfigClusterResolver . Jest podobny do TF_CONFIG w MultiWorkerMirroredStrategy ale ma dodatkowe zastrzeżenia.

W TF 1 ParameterServerStrategy jest dostępny tylko z estymatorem za tf.compat.v1.distribute.experimental.ParameterServerStrategy symbolu tf.compat.v1.distribute.experimental.ParameterServerStrategy .

CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy również przeprowadza trening synchroniczny. Zmienne nie są dublowane, zamiast tego są umieszczane na procesorze, a operacje są replikowane na wszystkich lokalnych GPU. Jeśli jest tylko jeden GPU, wszystkie zmienne i operacje zostaną umieszczone na tym GPU.

Utwórz instancję CentralStorageStrategy przez:

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'

Spowoduje to utworzenie instancji CentralStorageStrategy która będzie używać wszystkich widocznych procesorów graficznych i procesorów. Aktualizacja zmiennych w replikach zostanie zagregowana przed zastosowaniem do zmiennych.

Inne strategie

Oprócz powyższych strategii istnieją dwie inne strategie, które mogą być przydatne do tworzenia prototypów i debugowania podczas korzystania z interfejsów API tf.distribute .

Strategia domyślna

Strategia domyślna to strategia dystrybucji, która jest obecna, gdy w zakresie nie ma jawnej strategii dystrybucji. Implementuje interfejs tf.distribute.Strategy , ale jest tranzytem i nie zapewnia rzeczywistej dystrybucji. Na przykład strategy.run(fn) wywoła po prostu fn . Kod napisany przy użyciu tej strategii powinien zachowywać się dokładnie tak, jak kod napisany bez żadnej strategii. Możesz myśleć o tym jako o strategii „no-op”.

Strategia domyślna jest singletonem - i nie można tworzyć jej więcej. Można go uzyskać za pomocą tf.distribute.get_strategy() poza zakresem dowolnej jawnej strategii (ten sam interfejs API, którego można użyć do pobrania bieżącej strategii wewnątrz zakresu jawnej strategii).

default_strategy = tf.distribute.get_strategy()

Ta strategia służy dwóm głównym celom:

  • Pozwala bezwarunkowo pisać kod biblioteki obsługującej dystrybucję. Na przykład, w tf.optimizer s może użyć tf.distribute.get_strategy() i użyć tej strategii do redukcji gradientów - zawsze zwróci obiekt strategii, na którym możemy wywołać reduktor 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
  • Podobnie jak kod biblioteki, może być używany do pisania programów użytkowników końcowych do pracy ze strategią dystrybucji i bez niej, bez konieczności stosowania logiki warunkowej. Przykładowy fragment kodu ilustrujący to:
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.))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

OneDeviceStrategy

tf.distribute.OneDeviceStrategy to strategia umieszczania wszystkich zmiennych i obliczeń na jednym określonym urządzeniu.

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

Strategia ta różni się od strategii domyślnej pod wieloma względami. W strategii domyślnej logika umieszczania zmiennych pozostaje niezmieniona w porównaniu z uruchomionym TensorFlow bez żadnej strategii dystrybucji. Ale podczas korzystania z OneDeviceStrategy wszystkie zmienne utworzone w jego zakresie są jawnie umieszczane na określonym urządzeniu. Co więcej, wszelkie funkcje wywoływane przez OneDeviceStrategy.run również zostaną umieszczone na określonym urządzeniu.

Dane wejściowe dystrybuowane za pomocą tej strategii zostaną wstępnie pobrane do określonego urządzenia. W domyślnej strategii nie ma dystrybucji danych wejściowych.

Podobnie jak w przypadku strategii domyślnej, strategia ta może również służyć do testowania kodu przed przejściem na inne strategie, które faktycznie są dystrybuowane do wielu urządzeń / maszyn. Spowoduje to ćwiczenie mechanizmu strategii dystrybucji w nieco większym stopniu niż strategia domyślna, ale nie w pełnym zakresie, jak użycie MirroredStrategy lub TPUStrategy itp. Jeśli chcesz, aby kod zachowywał się tak, jakby nie miał strategii, użyj strategii domyślnej.

Jak dotąd widzieliście różne dostępne strategie i sposoby ich realizacji. W następnych kilku sekcjach przedstawiono różne sposoby wykorzystania ich do dystrybucji treningu.

Korzystanie tf.distribute.Strategy z tf.keras.Model.fit

tf.distribute.Strategy jest zintegrowana z tf.keras które jest implementacją specyfikacji API Keras TensorFlow. tf.keras to interfejs API wysokiego poziomu do tworzenia i trenowania modeli. Poprzez integrację z tf.keras sprawiliśmy, że bezproblemowo rozpowszechniamy szkolenie napisane w ramach szkoleniowych Keras przy użyciu model.fit .

Oto, co musisz zmienić w swoim kodzie:

  1. Utwórz instancję odpowiedniego tf.distribute.Strategy .
  2. Przenieś tworzenie modelu Keras, optymalizatora i metryk do strategy.scope . strategy.scope .

Obsługujemy wszystkie typy modeli Keras - sekwencyjne, funkcjonalne i podklasowe.

Oto fragment kodu, który to robi dla bardzo prostego modelu Keras z jedną gęstą warstwą:

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',).

W tym przykładzie zastosowano MirroredStrategy więc możesz uruchomić to na maszynie z wieloma GPU. strategy.scope() wskazuje Kerasowi, której strategii użyć do dystrybucji szkolenia. Tworzenie modeli / optymalizatorów / metryk w tym zakresie pozwala nam tworzyć zmienne rozproszone zamiast zwykłych zmiennych. Po skonfigurowaniu możesz dopasować model tak, jak zwykle. MirroredStrategy zajmuje się replikacją treningu modelu na dostępnych procesorach graficznych, agregowaniem gradientów i nie tylko.

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.9810
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.4336
10/10 [==============================] - 1s 2ms/step - loss: 0.2299

0.22988776862621307

Tutaj pliktf.data.Dataset zapewnia dane wejściowe dotyczące szkolenia i oceny. Możesz także użyć tablic 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.1636
Epoch 2/2
10/10 [==============================] - 0s 2ms/step - loss: 0.0723

<tensorflow.python.keras.callbacks.History at 0x7f89057a2470>

W obu przypadkach (zbiór danych lub numpy) każda partia danych wejściowych jest równo dzielona między wiele replik. Na przykład, jeśli używasz MirroredStrategy z 2 GPU, każda partia o rozmiarze 10 zostanie podzielona między 2 GPU, a każdy z nich otrzyma 5 przykładów wejściowych w każdym kroku. Każda epoka będzie trenować szybciej, gdy dodasz więcej GPU. Zwykle chciałbyś zwiększyć rozmiar partii, dodając więcej akceleratorów, aby efektywnie wykorzystać dodatkową moc obliczeniową. Będziesz także musiał ponownie dostroić współczynnik uczenia się, w zależności od modelu. Możesz użyć strategy.num_replicas_in_sync aby uzyskać liczbę replik.

# 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]

Co jest teraz obsługiwane?

Szkoleniowy interfejs API MirroredStrategy Strategia TPUS MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageStrategy
Interfejsy API Keras Utrzymany Utrzymany Wsparcie eksperymentalne Wsparcie eksperymentalne Wsparcie eksperymentalne

Przykłady i samouczki

Oto lista samouczków i przykładów, które ilustrują powyższą integrację od końca do końca z Kerasem:

  1. Samouczek do trenowania MNIST z MirroredStrategy .
  2. Samouczek do trenowania MNIST przy użyciu MultiWorkerMirroredStrategy .
  3. Przewodnik po treningu MNIST z wykorzystaniem TPUStrategy .
  4. Samouczek dotyczący szkolenia serwera parametrów w TensorFlow 2 z ParameterServerStrategy .
  5. TensorFlow model Garden repozytorium zawierającego zbiory state-of-the-art modeli realizowanych przy użyciu różnych strategii.

Używanie tf.distribute.Strategy z niestandardowymi pętlami szkoleniowymi

Jak widzieliśmy, używanie tf.distribute.Strategy z Keras model.fit wymaga zmiany tylko kilku wierszy kodu. Przy odrobinie wysiłku możesz również użyć tf.distribute.Strategy z niestandardowymi pętlami treningowymi.

Jeśli potrzebujesz większej elastyczności i kontroli nad pętlami treningowymi niż jest to możliwe w przypadku Estimatora lub Keras, możesz napisać niestandardowe pętle treningowe. Na przykład, korzystając z GAN, możesz chcieć wykonać inną liczbę kroków generatora lub dyskryminatora w każdej rundzie. Podobnie, ramy wysokiego poziomu nie nadają się zbytnio do szkolenia ze wzmocnieniem.

Klasy tf.distribute.Strategy zapewniają podstawowy zestaw metod do obsługi niestandardowych pętli szkoleniowych. Korzystanie z nich może początkowo wymagać niewielkiej restrukturyzacji kodu, ale gdy to zrobisz, powinieneś być w stanie przełączać się między GPU, TPU i wieloma maszynami, po prostu zmieniając instancję strategii.

Tutaj pokażemy krótki fragment ilustrujący ten przypadek użycia dla prostego przykładu szkoleniowego przy użyciu tego samego modelu Keras jak poprzednio.

Najpierw utwórz model i optymalizator w zakresie strategii. Dzięki temu wszelkie zmienne utworzone za pomocą modelu i optymalizatora są zmiennymi lustrzanymi.

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

Następnie utwórz wejściowy zestaw danych i wywołaj tf.distribute.Strategy.experimental_distribute_dataset celu dystrybucji zestawu danych w oparciu o strategię.

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

Następnie określ jeden etap szkolenia. Użyjtf.GradientTape aby obliczyć gradienty i optymalizatora, aby zastosować te gradienty do zaktualizowania zmiennych naszego modelu. Aby rozpowszechnić ten krok szkoleniowy, umieść go w funkcji train_step i przekaż do tf.distrbute.Strategy.run wraz z danymi wejściowymi zestawu danych uzyskanymi z dist_dataset utworzonego wcześniej:

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)

Kilka innych rzeczy, o których należy pamiętać w powyższym kodzie:

  1. Użył tf.nn.compute_average_loss do obliczenia straty. tf.nn.compute_average_loss sumuje straty na przykład i dzieli sumę przez global_batch_size. Jest to ważne, ponieważ później, po obliczeniu gradientów na każdej replice, są one agregowane w replikach poprzez zsumowanie .
  2. tf.distribute.Strategy.reduce interfejsu API tf.distribute.Strategy.reduce do agregowania wyników zwróconych przez tf.distribute.Strategy.run . tf.distribute.Strategy.run zwraca wyniki z każdej lokalnej repliki w strategii i istnieje wiele sposobów wykorzystania tego wyniku. Możesz je reduce , aby uzyskać zagregowaną wartość. Możesz również wykonać tf.distribute.Strategy.experimental_local_results aby uzyskać listę wartości zawartych w wyniku, po jednej na lokalną replikę.
  3. Gdy apply_gradients jest wywoływana w zakresie strategii dystrybucji, jego zachowanie jest modyfikowane. W szczególności przed zastosowaniem gradientów w każdej równoległej instancji podczas treningu synchronicznego wykonuje sumaryczne repliki wszystkich gradientów.

Wreszcie, po zdefiniowaniu kroku szkoleniowego, możemy iterować po dist_dataset i uruchomić szkolenie w pętli:

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))
tf.Tensor(0.88850546, shape=(), dtype=float32)
tf.Tensor(0.8815902, shape=(), dtype=float32)
tf.Tensor(0.87474185, shape=(), dtype=float32)
tf.Tensor(0.8679599, shape=(), dtype=float32)
tf.Tensor(0.86124384, shape=(), dtype=float32)
tf.Tensor(0.8545932, shape=(), dtype=float32)
tf.Tensor(0.8480073, shape=(), dtype=float32)
tf.Tensor(0.8414857, shape=(), dtype=float32)
tf.Tensor(0.8350279, shape=(), dtype=float32)
tf.Tensor(0.8286335, shape=(), dtype=float32)
tf.Tensor(0.8223018, shape=(), dtype=float32)
tf.Tensor(0.81603223, shape=(), dtype=float32)
tf.Tensor(0.8098244, shape=(), dtype=float32)
tf.Tensor(0.8036777, shape=(), dtype=float32)
tf.Tensor(0.7975916, shape=(), dtype=float32)
tf.Tensor(0.79156566, shape=(), dtype=float32)
tf.Tensor(0.78559923, shape=(), dtype=float32)
tf.Tensor(0.77969193, shape=(), dtype=float32)
tf.Tensor(0.773843, shape=(), dtype=float32)
tf.Tensor(0.7680521, shape=(), dtype=float32)

W powyższym przykładzie dist_dataset po dist_dataset aby zapewnić dane wejściowe do treningu. tf.distribute.Strategy.make_experimental_numpy_dataset również tf.distribute.Strategy.make_experimental_numpy_dataset danych tf.distribute.Strategy.make_experimental_numpy_dataset do obsługi danych wejściowych numpy. Możesz użyć tego interfejsu API do utworzenia zestawu danych przed wywołaniem tf.distribute.Strategy.experimental_distribute_dataset .

Innym sposobem iterowania danych jest jawne użycie iteratorów. Możesz to zrobić, jeśli chcesz wykonać określoną liczbę kroków, a nie iterować po całym zestawie danych. Powyższy iteracji będzie teraz być modyfikowana najpierw utworzyć iterator a następnie jawnie wywołać next na to, aby uzyskać dane wejściowe.

iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))
tf.Tensor(0.7623187, shape=(), dtype=float32)
tf.Tensor(0.7566423, shape=(), dtype=float32)
tf.Tensor(0.7510221, shape=(), dtype=float32)
tf.Tensor(0.7454578, shape=(), dtype=float32)
tf.Tensor(0.739949, shape=(), dtype=float32)
tf.Tensor(0.7344949, shape=(), dtype=float32)
tf.Tensor(0.72909516, shape=(), dtype=float32)
tf.Tensor(0.7237492, shape=(), dtype=float32)
tf.Tensor(0.7184567, shape=(), dtype=float32)
tf.Tensor(0.7132167, shape=(), dtype=float32)

Obejmuje to najprostszy przypadek użycia interfejsu API tf.distribute.Strategy do dystrybucji niestandardowych pętli szkoleniowych.

Co jest teraz obsługiwane?

Szkoleniowy interfejs API MirroredStrategy Strategia TPUS MultiWorkerMirroredStrategy ParameterServerStrategy CentralStorageStrategy
Niestandardowa pętla treningowa Utrzymany Utrzymany Wsparcie eksperymentalne Wsparcie eksperymentalne Wsparcie eksperymentalne

Przykłady i samouczki

Oto kilka przykładów użycia strategii dystrybucji z niestandardowymi pętlami szkoleniowymi:

  1. Samouczek do trenowania MNIST przy użyciu MirroredStrategy .
  2. Przewodnik po szkoleniu MNIST z wykorzystaniem TPUStrategy .
  3. TensorFlow model Garden repozytorium zawierającego zbiory state-of-the-art modeli realizowanych przy użyciu różnych strategii.

Inne tematy

W tej sekcji omówiono kilka tematów, które dotyczą wielu przypadków użycia.

Konfigurowanie zmiennej środowiskowej TF_CONFIG

W przypadku szkolenia wielu pracowników, jak wspomniano wcześniej, należy ustawić TF_CONFIG środowiskową TF_CONFIG dla każdego TF_CONFIG binarnego działającego w klastrze. TF_CONFIG środowiskowa TF_CONFIG jest ciągiem JSON, który określa, jakie zadania tworzą klaster, ich adresy i rolę każdego zadania w klastrze. Repozytorium tensorflow / ekosystem udostępnia szablon Kubernetes, w którym ustawia TF_CONFIG dla zadań szkoleniowych.

Istnieją dwa składniki TF_CONFIG: klaster i zadanie. klaster dostarcza informacji o klastrze szkoleniowym, który jest dyktatem składającym się z różnych rodzajów zawodów, takich jak pracownik. W przypadku szkolenia dla wielu pracowników zwykle jeden pracownik przejmuje nieco większą odpowiedzialność, taką jak zapisywanie punktu kontrolnego i pisanie pliku podsumowania dla TensorBoard, oprócz tego, co robi zwykły pracownik. Taki pracownik jest nazywany „szefem” i jest zwyczajem, że pracownik o indeksie 0 jest mianowany jako główny pracownik (w rzeczywistości jest to sposób implementacji strategii tf.distribute.Strategy). zadanie z drugiej strony dostarcza informacji o bieżącym zadaniu. Pierwszy klaster składników jest taki sam dla wszystkich pracowników, a drugie zadanie składnika jest różne dla każdego pracownika i określa typ i indeks tego pracownika.

Jednym z przykładów TF_CONFIG jest:

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

Ten TF_CONFIG określa, że ​​w TF_CONFIG istnieją trzy TF_CONFIG i dwa zadania ps wraz z ich hostami i portami. Część „zadanie” określa rolę aktualnego zadania w klastrze, pracownik 1 (drugi pracownik). Prawidłowe role w klastrze to „szef”, „pracownik”, „ps” i „oceniający”. Zadanie "ps" nie powinno istnieć, z wyjątkiem użycia tf.distribute.experimental.ParameterServerStrategy .

Co dalej?

tf.distribute.Strategy jest aktywnie rozwijana. Wypróbuj i przekaż swoją opinię za pomocą problemów z GitHub .