TensorFlow-Code auf TPU-Pod-Slices ausführen
In diesem Dokument erfahren Sie, wie Sie mit TensorFlow eine Berechnung auf einem TPU-Pod ausführen. Führen Sie dazu die folgenden Schritte aus:
- TPU-Pod-Slice mit TensorFlow-Software erstellen
- SSH-Verbindung zur TPU-VM herstellen
- Beispielskript erstellen und ausführen
Die TPU-VM benötigt ein Dienstkonto
nach Berechtigungen zum Aufrufen der Cloud TPU API. Standardmäßig verwendet Ihre TPU-VM das Standarddienstkonto der Compute Engine, das alle erforderlichen Cloud TPU-Berechtigungen enthält. Wenn Sie Ihre eigene
Dienstkonto, dem Sie den TPU Viewer hinzufügen müssen
Rolle zu Ihrem Dienstkonto. Weitere Informationen zu Google Cloud-Rollen finden Sie unter Informationen zu Rollen.
Sie können Ihr eigenes Dienstkonto mit dem Flag --service-account
angeben, wenn
die TPU-VM erstellen.
Umgebung einrichten
Führen Sie in der Cloud Shell den folgenden Befehl aus, um sicherzustellen, dass Sie die aktuelle Version von
gcloud
ausführen:$ gcloud components update
Wenn Sie
gcloud
installieren möchten, verwenden Sie den folgenden Befehl:$ sudo apt install -y google-cloud-sdk
Erstellen Sie einige Umgebungsvariablen:
$ export PROJECT_ID=project-id $ export TPU_NAME=tpu-name $ export ZONE=europe-west4-a $ export RUNTIME_VERSION=tpu-vm-tf-2.17.0-pod-pjrt $ export ACCELERATOR_TYPE=v3-32
TPU v3-32-Pod-Slice mit TensorFlow-Laufzeit erstellen
$ gcloud compute tpus tpu-vm create ${TPU_NAME}} \ --zone=${ZONE} \ --accelerator-type=${ACCELERATOR_TYPE} \ --version=${RUNTIME_VERSION}
Beschreibung der Befehls-Flags
zone
- Die Zone, in der Sie die 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 jede TPU-Version finden Sie unter TPU-Versionen.
version
- Die Softwareversion von Cloud TPU.
Über SSH eine Verbindung zur Cloud TPU-VM herstellen
$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --zone=${ZONE}
Beispielskript erstellen und ausführen
Legen Sie die folgenden Umgebungsvariablen fest:
(vm)$ export TPU_NAME=tpu-name (vm)$ export TPU_LOAD_LIBRARY=0
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)
Führen Sie dieses Skript mit dem folgenden Befehl aus:
(vm)$ python3 tpu-test.py
Dieses Skript führt eine Berechnung für jeden TensorCore eines TPU-Pod-Slice aus. 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.
Trennen Sie die Verbindung zur Compute Engine:
(vm)$ exit
Löschen Sie Ihre Cloud TPU.
$ gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --zone=${ZONE}
Ü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=${ZONE}