Berechnung mit TensorFlow 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. Eine ausführlichere Anleitung, die zeigt, wie Sie ein Modell auf einer Cloud TPU trainieren, finden Sie in einer Cloud TPU-Anleitung.

Hinweise

Bevor Sie dieser Kurzanleitung folgen, müssen Sie ein Google Cloud-Konto erstellen, die Google Cloud CLI installieren und den Befehl gcloud 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.18.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 den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.
version
Die Version der TPU-Laufzeit In der in dieser Kurzanleitung verwendeten Version ist TensorFlow bereits vorinstalliert.
project
Der Name des Google Cloud CLI-Projekts, in dem Sie die 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 zu Ihrer 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

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

    import tensorflow as tf

print(f"TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores")

  2. Führen Sie das Skript aus:

    (vm)$ python3 tpu-count.py

    Die Ausgabe des Scripts zeigt die Anzahl der TPU-Kerne an, die für die TPU-VM verfügbar sind:

    TensorFlow can access 8 TPU cores

Einfache Berechnung mit TensorFlow ausführen

Nachdem Sie eine Verbindung zur TPU-VM hergestellt haben, legen Sie die folgende Umgebungsvariable fest.

  (vm)$ export TPU_NAME=local

Wenn Sie beim Erstellen Ihrer TPU den Parameter --version auf eine Version mit der Endung -pjrt festlegen, müssen Sie die folgenden Umgebungsvariablen festlegen, 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:

  1. 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.

  2. Löschen Sie Ihre Cloud TPU.

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

  3. Prü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: