Diese Seite wurde von der Cloud Translation API übersetzt.

TensorFlow-Code auf TPU-Pod-Segmenten ausführen

In diesem Dokument erfahren Sie, wie Sie mit TensorFlow auf einem TPU-Pod eine einfache Berechnung ausführen. Sie werden die folgenden Schritte ausführen:

TPU-Pod-Slice mit TensorFlow-Software erstellen
Über SSH eine Verbindung zur TPU-VM herstellen
Ein einfaches Skript erstellen und ausführen

Die TPU-VM benötigt Dienstkonten für Berechtigungen zum Aufrufen der Cloud TPU API. Standardmäßig verwendet die TPU-VM das standardmäßige Compute Engine-Dienstkonto, das alle erforderlichen Cloud TPU-Berechtigungen enthält. Wenn Sie Ihr eigenes Dienstkonto verwenden, müssen Sie Ihrem Dienstkonto die Rolle TPU-Betrachter hinzufügen. Weitere Informationen zu Google Cloud-Rollen finden Sie unter Informationen zu Rollen. Sie können beim Erstellen Ihrer TPU-VM mit dem Flag --service-account Ihr eigenes Dienstkonto angeben.

v3-32 TPU-Pod-Slice mit TensorFlow-Laufzeit erstellen

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

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 Version der Cloud TPU-Software. Standardmäßig wird die neueste TensorFlow-Softwareversion verwendet.

Verbindung zur Cloud TPU-VM über SSH herstellen

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

Einfaches Berechnungsskript 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 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),
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 europe-west4-a

Ü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 europe-west4-a
```