Kurzanleitung: Berechnungen 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 auf einer Cloud TPU ausführen. Eine ausführlichere Anleitung zum Trainieren eines Modells auf einer Cloud TPU finden Sie in den Cloud TPU-Anleitungen.

Hinweis

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 eine Compute Engine Cloud TPU mit dem Befehl gcloud. 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 gcloud-Referenz.

TPU-VM

$ gcloud compute tpus tpu-vm create tpu-name \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.11.0

Beschreibung der Befehls-Flags

zone
Die Zone, in der Sie Ihre Cloud TPU erstellen möchten.
accelerator-type
Der Typ der zu erstellenden Cloud TPU.
version
Die Cloud TPU-Softwareversion. Standardmäßig wird die neueste TensorFlow-Softwareversion verwendet.

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.11.0 \

Beschreibung der Befehls-Flags

project
Ihre GCP-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 der VM, die mit dem Befehl gcloud erstellt wurde, in GB.
machine-type
Der Maschinentyp der zu erstellenden Compute Engine-VM.
tf-version
Die Version von TensorFlow gcloud wird auf der VM installiert.
accelerator-type
Der Typ der zu erstellenden Cloud TPU.

Verbindung zur Cloud TPU-VM herstellen

Wenn Sie TPU-VMs verwenden, müssen Sie explizit eine SSH-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

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

(vm)$ export TPU_NAME=local

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.

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 Skript führt eine einfache Berechnung auf jedem TensorCore einer TPU 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

So 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: