TensorFlow-Code auf TPU-Pod-Slices ausführen

In diesem Dokument erfahren Sie, wie Sie mit TensorFlow eine einfache Berechnung auf einem TPU-Pod ausführen. Sie werden die folgenden Schritte ausführen:

1. TPU-Pod-Slice mit TensorFlow-Software erstellen
2. Über SSH eine Verbindung zur TPU-VM herstellen
3. Einfaches Skript erstellen und ausführen

Die TPU-VM benötigt Dienstkonten für Berechtigungen zum Aufrufen der Cloud TPU API. Standardmäßig verwendet Ihre TPU-VM das Compute Engine-Standarddienstkonto, das alle erforderlichen Cloud TPU-Berechtigungen enthält. Wenn Sie ein eigenes Dienstkonto verwenden, müssen Sie Ihrem Dienstkonto die Rolle TPU Viewer hinzufügen. Weitere Informationen zu Google Cloud-Rollen finden Sie unter Informationen zu Rollen. Mit dem Flag --service-account können Sie beim Erstellen Ihrer TPU-VM ein eigenes Dienstkonto angeben.

v3-32-TPU-Pod-Slice mit TensorFlow-Laufzeit erstellen

Hinweis: Beim Starten der TPU-VM können Sie entweder die im Befehl angegebene TPU-Softwareversion verwenden oder eine andere unterstützte TPU-Softwareversion unter Cloud TPU-Softwareversionen angeben.

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.16.1-pod-pjrt

Beschreibung der Befehls-Flags

zone

Die Zone, in der Sie Ihre Cloud TPU erstellen möchten.

accelerator-type

Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.

version

Die Cloud TPU-Softwareversion.

Über SSH eine Verbindung zur Cloud TPU-VM herstellen

$ gcloud compute tpus tpu-vm ssh tpu-name \
      --zone europe-west4-a

Einfaches Berechnungsskript erstellen und ausführen

1. Legen Sie die folgenden Umgebungsvariablen fest:

   (vm)$ export TPU_NAME=tpu-name
   (vm)$ export TPU_LOAD_LIBRARY=0
   

2. Erstellen Sie im aktuellen Verzeichnis eine Datei mit dem Namen tpu-test.py, kopieren Sie das folgende Skript und fügen Sie es in diese Datei ein.

   import tensorflow as tf
   print("Tensorflow version " + tf.__version__)
   
   cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
   print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])
   
   tf.config.experimental_connect_to_cluster(cluster_resolver)
   tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
   strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)
   
   @tf.function
   def add_fn(x,y):
   z = x + y
   return z
   
   x = tf.constant(1.)
   y = tf.constant(1.)
   z = strategy.run(add_fn, args=(x,y))
   print(z)
   

3. Führen Sie dieses Skript mit dem folgenden Befehl aus:

   (vm)$ python3 tpu-test.py

   Dieses Script führt auf jedem TensorCore einer TPU eine einfache Berechnung durch. Die Ausgabe sollte in etwa so aussehen:

   PerReplica:{
   0: tf.Tensor(2.0, shape=(), dtype=float32),
   1: tf.Tensor(2.0, shape=(), dtype=float32),
   2: tf.Tensor(2.0, shape=(), dtype=float32),
   3: tf.Tensor(2.0, shape=(), dtype=float32),
   4: tf.Tensor(2.0, shape=(), dtype=float32),
   5: tf.Tensor(2.0, shape=(), dtype=float32),
   6: tf.Tensor(2.0, shape=(), dtype=float32),
   7: tf.Tensor(2.0, shape=(), dtype=float32),
   8: tf.Tensor(2.0, shape=(), dtype=float32),
   9: tf.Tensor(2.0, shape=(), dtype=float32),
   10: tf.Tensor(2.0, shape=(), dtype=float32),
   11: tf.Tensor(2.0, shape=(), dtype=float32),
   12: tf.Tensor(2.0, shape=(), dtype=float32),
   13: tf.Tensor(2.0, shape=(), dtype=float32),
   14: tf.Tensor(2.0, shape=(), dtype=float32),
   15: tf.Tensor(2.0, shape=(), dtype=float32),
   16: tf.Tensor(2.0, shape=(), dtype=float32),
   17: tf.Tensor(2.0, shape=(), dtype=float32),
   18: tf.Tensor(2.0, shape=(), dtype=float32),
   19: tf.Tensor(2.0, shape=(), dtype=float32),
   20: tf.Tensor(2.0, shape=(), dtype=float32),
   21: tf.Tensor(2.0, shape=(), dtype=float32),
   22: tf.Tensor(2.0, shape=(), dtype=float32),
   23: tf.Tensor(2.0, shape=(), dtype=float32),
   24: tf.Tensor(2.0, shape=(), dtype=float32),
   25: tf.Tensor(2.0, shape=(), dtype=float32),
   26: tf.Tensor(2.0, shape=(), dtype=float32),
   27: tf.Tensor(2.0, shape=(), dtype=float32),
   28: tf.Tensor(2.0, shape=(), dtype=float32),
   29: tf.Tensor(2.0, shape=(), dtype=float32),
   30: tf.Tensor(2.0, shape=(), dtype=float32),
   31: tf.Tensor(2.0, shape=(), dtype=float32)
   }
   

Bereinigen

Wenn Sie mit Ihrer TPU-VM fertig sind, führen Sie die folgenden Schritte aus, um Ihre Ressourcen zu bereinigen.

1. Trennen Sie die Verbindung zur Compute Engine:

   (vm)$ exit
   

2. Löschen Sie Ihre Cloud TPU.

   $ gcloud compute tpus tpu-vm delete tpu-name \
     --zone europe-west4-a
   

3. Überprüfen Sie mit dem folgenden Befehl, ob die Ressourcen gelöscht wurden. Achten Sie darauf, dass Ihre TPU nicht mehr aufgeführt wird. Der Löschvorgang kann einige Minuten dauern.

   $ gcloud compute tpus tpu-vm list \
     --zone europe-west4-a