Dzień Społeczności ML jest 9 listopada! Dołącz do nas na aktualizacje z TensorFlow Jax i więcej Dowiedz się więcej

Migracja z TPUEstimator do TPUStrategy

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

Ten poradnik pokazuje, jak przenieść swoje obieg działa na TPU z TensorFlow 1 za TPUEstimator API TensorFlow 2 za TPUStrategy API.

Na koniec-koniec TensorFlow 2 przykłady, sprawdź Zastosowanie TPU przewodnik-mianowicie Klasyfikacja na TPU sekcja-a Rozwiązywanie zadań z wykorzystaniem BERT na klej TPU tutoriala. Można również znaleźć Distributed szkoleniowy przewodnik powinien być przydatny, która obejmuje wszystkie strategie dystrybucji TensorFlow, w tym TPUStrategy .

Ustawiać

Zacznij od importów i prostego zestawu danych do celów demonstracyjnych:

import tensorflow as tf
import tensorflow.compat.v1 as tf1
features = [[1., 1.5]]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_labels = [[0.8]]

TensorFlow 1: Napęd modelu na TPU za pomocą TPUEstimator

Ta część przewodnika pokazuje, jak przeprowadzić szkolenia i oceny z tf.compat.v1.estimator.tpu.TPUEstimator w TensorFlow 1.

Aby użyć TPUEstimator , najpierw zdefiniować kilka funkcji: funkcję wejściową dla danych treningowych, funkcję wejściową ocena dla danych ewaluacyjnych oraz funkcję modelu, który opowiada TPUEstimator jak op szkolenie jest zdefiniowana z funkcjami i etykiety:

def _input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((features, labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _eval_input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((eval_features, eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.tpu.TPUEstimatorSpec(mode, loss=loss, train_op=train_op)

Z zdefiniowane te funkcje, należy utworzyć tf.distribute.cluster_resolver.TPUClusterResolver który dostarcza informacji klastra oraz tf.compat.v1.estimator.tpu.RunConfig obiekt. Wraz z funkcją modelu zdefiniowanego, można teraz utworzyć TPUEstimator . Tutaj uprościsz przepływ, pomijając oszczędności w punktach kontrolnych. Następnie można określić wielkość wsadu zarówno szkolenia i oceny dla TPUEstimator .

cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')
print("All devices: ", tf1.config.list_logical_devices('TPU'))
tpu_config = tf1.estimator.tpu.TPUConfig(iterations_per_loop=10)
config = tf1.estimator.tpu.RunConfig(
    cluster=cluster_resolver,
    save_checkpoints_steps=None,
    tpu_config=tpu_config)
estimator = tf1.estimator.tpu.TPUEstimator(
    model_fn=_model_fn,
    config=config,
    train_batch_size=8,
    eval_batch_size=8)

Zadzwoń TPUEstimator.train rozpocząć szkolenie modelu:

estimator.train(_input_fn, steps=1)

Następnie zadzwonić TPUEstimator.evaluate do oceny modelu przy użyciu danych z oceny:

estimator.evaluate(_eval_input_fn, steps=1)

TensorFlow 2: Prowadź model na TPU za pomocą Keras Model.fit i TPUStrategy

W TensorFlow 2, trenować na pracowników TPU, użyj tf.distribute.TPUStrategy wraz z API Keras definicja modelu i szkolenia / ewaluacji. (Patrz Zastosowanie TPU przewodnika więcej przykładami treningu z Keras Model.fit oraz pętlę własnych treningów (z tf.function i tf.GradientTape )).

Ponieważ trzeba wykonać jakąś pracę inicjalizacji połączyć się ze zdalnym klastrze i zainicjować robotników TPU, zacząć od stworzenia TPUClusterResolver dostarczenie informacji klastra i podłączyć do klastra. (Więcej informacji w sekcji inicjującej TPU z Stosowanie TPU przewodnika).

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

Następnie, gdy dane są przygotowane, można utworzyć TPUStrategy , zdefiniowanie modelu, metryki, oraz optymalizator pod zakresem niniejszej strategii.

Aby osiągnąć porównywalną prędkość szkoleniowy z TPUStrategy , należy upewnić się, aby wybrać numer dla steps_per_execution w Model.compile ponieważ określa liczbę partii do uruchomienia podczas każdego tf.function rozmowy i ma decydujące znaczenie dla wydajności. Argument ten jest podobny do iterations_per_loop wykorzystywane w TPUEstimator . Jeśli używasz pętle zwyczaj szkoleniowe, należy upewnić się, że wiele etapów są prowadzone w ramach tf.function -ED funkcji szkoleniowej. Przejdź do poprawy wyników z wielu etapów wewnątrz sekcji tf.function o Wykorzystanie TPU przewodnik dla dalszych informacji.

tf.distribute.TPUStrategy może wspierać ograniczone kształty dynamiczne, co jest tak, że górna granica dynamicznego kształtu obliczeń można wywnioskować. Jednak kształty dynamiczne mogą powodować pewne narzuty na wydajność w porównaniu z kształtami statycznymi. Dlatego ogólnie zaleca się, aby kształty wejściowe były statyczne, jeśli to możliwe, zwłaszcza podczas treningu. Jeden wspólny op która zwraca kształt dynamicznej jest tf.data.Dataset.batch(batch_size) , ponieważ liczba próbek pozostających w strumieniu może być mniejsza od wielkości partii. Dlatego też, gdy szkolenie w TPU, należy użyć tf.data.Dataset.batch(..., drop_remainder=True) dla uzyskania najlepszej wydajności treningowej.

dataset = tf.data.Dataset.from_tensor_slices(
    (features, labels)).shuffle(10).repeat().batch(
        8, drop_remainder=True).prefetch(2)
eval_dataset = tf.data.Dataset.from_tensor_slices(
    (eval_features, eval_labels)).batch(1, drop_remainder=True)

strategy = tf.distribute.TPUStrategy(cluster_resolver)
with strategy.scope():
  model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)
  model.compile(optimizer, "mse", steps_per_execution=10)

Dzięki temu jesteś gotowy do trenowania modelu za pomocą uczącego zestawu danych:

model.fit(dataset, epochs=5, steps_per_epoch=10)

Na koniec oceń model przy użyciu zestawu danych ewaluacyjnych:

model.evaluate(eval_dataset, return_dict=True)

Następne kroki

Aby dowiedzieć się więcej o TPUStrategy w TensorFlow 2, należy rozważyć następujące zasoby:

Aby dowiedzieć się więcej o dostosowywaniu treningu, zapoznaj się z:

TPU-Google specjalizowane układy ASIC do maszynowego uczenia są dostępne za pośrednictwem Google Colab , z TPU Badań chmurze , a chmura TPU .