Berechnung auf einer Cloud TPU-VM mit TensorFlow ausführen

Berechnung auf einer Cloud TPU-VM mit TensorFlow ausführen

In dieser Kurzanleitung erfahren Sie, wie Sie eine Cloud TPU erstellen, TensorFlow installieren und eine einfache Berechnung in einer Cloud TPU ausführen. Eine ausführliche Anleitung zum Trainieren eines Modells auf einer Cloud TPU finden Sie in den Cloud TPU-Anleitungen.

Hinweise

Wichtig: Sie können diese Kurzanleitung entweder mit der TPU-VM- oder mit der TPU-Knotenkonfiguration verwenden. Die beiden VM-Architekturen werden unter Systemarchitektur beschrieben. Welche gcloud-Befehle Sie verwenden, hängt von der verwendeten TPU-Konfiguration ab. In dieser Kurzanleitung wird jeder gcloud-Befehl in einem eigenen Tab angezeigt. Wählen Sie den Tab für die TPU-Konfiguration aus, die Sie verwenden möchten. Auf der Webseite wird der entsprechende gcloud-Befehl angezeigt. Sofern Sie nicht unbedingt TPU-Knoten nutzen müssen, empfehlen wir die Verwendung von TPU-VMs.

Bevor Sie dieser Kurzanleitung folgen, müssen Sie ein Google Cloud Platform-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 oder -Knoten mit gcloud erstellen

Starten Sie mit dem Befehl gcloud eine Compute Engine Cloud TPU. Der verwendete Befehl hängt davon ab, ob Sie eine TPU-VM oder einen TPU-Knoten verwenden. Weitere Informationen zu beiden VM-Architekturen finden Sie unter Systemarchitektur. Weitere Informationen zum Befehl gcloud finden Sie in der Referenz zu gcloud.

TPU-VM

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

$ gcloud compute tpus tpu-vm create tpu-name \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.16.1-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.

TPU-Knoten

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.12.0 \

Beschreibung der Befehls-Flags

project

Ihre Google Cloud-Projekt-ID

name

Der Name der zu erstellenden Cloud TPU.

zone

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

disk-size

Die Größe der Festplatte in GB der mit dem Befehl gcloud erstellten VM.

machine-type

Der Maschinentyp der Compute Engine-VM, die erstellt werden soll.

tf-version

Die Version von TensorFlow gcloud wird auf der VM installiert. Weitere Informationen finden Sie unter

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.

Hinweis: Wenn Sie mehrere Projekte haben, müssen Sie die Projekt-ID mit dem Flag --project angeben.

Verbindung zur Cloud TPU-VM herstellen

Wenn Sie TPU-VMs verwenden, müssen Sie über SSH explizit eine Verbindung zu Ihrer TPU-VM herstellen. Wenn Sie TPU-Knoten verwenden, sollten Sie automatisch eine SSH-Verbindung zu Ihrer Compute Engine-VM herstellen. Wenn Sie nicht automatisch verbunden werden, verwenden Sie den folgenden Befehl.

TPU-VM

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

TPU-Knoten

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

Einfaches Beispiel mit TensorFlow ausführen

TPU-VM

Hinweis: In diesem Beispiel wird gezeigt, wie Code auf einer einzelnen TPU ausgeführt wird (z. B. v2-8 oder v3-8). Informationen zum Ausführen von Code auf einem größeren TPU-Slice oder Pod (z. B. v2-32+ oder v3-32+) finden Sie im Pod-Beispiel.

Sobald Sie mit der TPU-VM verbunden sind, legen Sie die folgende Umgebungsvariable fest.

  (vm)$ export TPU_NAME=local

Wenn Sie beim Erstellen Ihrer TPU den Parameter --version auf eine Version festlegen, die auf -pjrt endet, 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)

TPU-Knoten

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.

Hinweis: Übergeben Sie bei einzelnen TPUs, Slices und Pods den TPU-Namen an TPUClusterResolver().

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)

@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 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)
}

Bereinigen

Mit den folgenden Schritten vermeiden Sie, dass Ihrem Google Cloud-Konto die auf dieser Seite 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.

   TPU-VM

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

   TPU-Knoten

   $ gcloud compute tpus execution-groups 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.

   TPU-VM

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

   TPU-Knoten

   $ gcloud compute tpus execution-groups list --zone=europe-west4-a
   

Nächste Schritte

Weitere Informationen zu Cloud TPU finden Sie unter:

• TensorFlow-Code auf TPU-Pod-Slices ausführen
• TPUs verwalten
• Cloud TPU-Systemarchitektur