Mit TensorFlow eine Berechnung auf einer Cloud TPU-VM ausführen
In dieser Kurzanleitung erfahren Sie, wie Sie eine Cloud TPU erstellen und mit TensorFlow eine Berechnung auf einer Cloud TPU ausführen. Ein ausführlicheres Tutorial, in dem Informationen zum Trainieren eines Modells auf einer Cloud TPU finden Sie in einer der Cloud TPU-Anleitungen.
Hinweis
Bevor Sie dieser Kurzanleitung folgen, müssen Sie ein Google Cloud-Konto erstellen, die Google Cloud CLI installieren und den gcloud
-Befehl konfigurieren. Weitere Informationen finden Sie unter Konto und Cloud TPU-Projekt einrichten.
Cloud TPU-VM mit gcloud
erstellen
Erstellen Sie eine Cloud TPU mit dem Befehl gcloud
.
$ gcloud compute tpus tpu-vm create tpu-name
--zone=europe-west4-a
--accelerator-type=v3-8
--version=tpu-vm-tf-2.17.0-pjrt
--project=your-gcp-project-name
Beschreibung der Befehls-Flags
tpu-name
- Der Name der zu erstellenden Cloud TPU.
zone
- Die Zone, in der Sie die Cloud TPU erstellen möchten.
accelerator-type
- Mit dem Beschleunigertyp geben Sie 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 Version der TPU-Laufzeit Bei der in dieser Kurzanleitung verwendeten Version ist TensorFlow vorinstalliert.
project
- Der Name des Google Cloud CLI-Projekts, in dem Sie Ihre Cloud TPU erstellen.
Weitere Informationen zum Befehl gcloud
finden Sie in der gcloud
-Referenz.
Verbindung zur Cloud TPU-VM herstellen
Stellen Sie über SSH eine Verbindung zur TPU-VM her:
$ gcloud compute tpus tpu-vm ssh tpu-name
--zone europe-west4-a
--project=your-gcp-project-name
Prüfen, ob TensorFlow auf TPUs zugreifen kann
Erstellen Sie im aktuellen Verzeichnis eine Datei mit dem Namen
tpu-count.py
und kopieren Sie das folgende Skript ein.import tensorflow as tf print(f'TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores')
Führen Sie das Skript aus:
(vm)$ python3 tpu-count.py
Die Ausgabe des Skripts zeigt die Anzahl der für die TPU-VM verfügbaren TPU-Kerne an:
TensorFlow can access 8 TPU cores
Einfache Berechnung mit TensorFlow ausführen
Sobald Sie mit der TPU-VM verbunden sind, legen Sie die folgende Umgebungsvariable fest.
(vm)$ export TPU_NAME=local
Wenn Sie beim Erstellen der TPU den Parameter --version
auf eine Version festlegen, die mit
-pjrt
, legen Sie die folgenden Umgebungsvariablen fest, um die PJRT-Laufzeit zu aktivieren:
(vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
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__) @tf.function def add_fn(x,y): z = x + y return z cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver() tf.config.experimental_connect_to_cluster(cluster_resolver) tf.tpu.experimental.initialize_tpu_system(cluster_resolver) strategy = tf.distribute.TPUStrategy(cluster_resolver) 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 auf jedem TensorCore einer TPU eine Berechnung 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) }
Bereinigen
Mit den folgenden Schritten vermeiden Sie, dass Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen in Rechnung gestellt werden:
Trennen Sie die Verbindung zur Compute Engine-Instanz, sofern noch nicht geschehen:
(vm)$ exit
Die Eingabeaufforderung sollte nun
username@projectname
lauten und angeben, dass Sie sich in Cloud Shell befinden.Löschen Sie Ihre Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name
--zone=europe-west4-aPrüfen Sie, ob die Ressourcen gelöscht wurden. Führen Sie dazu
gcloud compute tpus tpu-vm list
aus. Der Löschvorgang kann einige Minuten dauern.$ gcloud compute tpus tpu-vm list --zone=europe-west4-a
Nächste Schritte
Weitere Informationen zu Cloud TPU finden Sie unter: