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

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

  1. TPU-Pod-Slice mit TensorFlow-Software erstellen
  2. SSH-Verbindung zur TPU-VM herstellen
  3. Beispielskript erstellen und ausführen

Die TPU-VM benötigt Dienstkonten, um die Cloud TPU API aufrufen zu können. Standardmäßig verwendet Ihre TPU-VM das Standarddienstkonto der Compute Engine, das alle erforderlichen Cloud TPU-Berechtigungen enthält. Wenn Sie ein eigenes Dienstkonto verwenden, müssen Sie ihm 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 ein eigenes Dienstkonto angeben.

Umgebung einrichten

  1. Führen Sie in der Cloud Shell den folgenden Befehl aus, um sicherzustellen, dass Sie die aktuelle Version von gcloud ausführen:

    $ gcloud components update

    Wenn Sie gcloud installieren möchten, verwenden Sie den folgenden Befehl:

    $ sudo apt install -y google-cloud-sdk

  2. Erstellen Sie einige Umgebungsvariablen:

    $ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=europe-west4-a
$ export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

TPU-Pod-Slice der Version 3.32 mit TensorFlow-Laufzeit erstellen

$ gcloud compute tpus tpu-vm create ${TPU_NAME}} \
  --zone=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE} \
  --version=${RUNTIME_VERSION}

Beschreibung der Befehls-Flags

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 Cloud TPU-Softwareversion.

Über SSH eine Verbindung zur Cloud TPU-VM herstellen

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --zone=${ZONE}

Beispielskript erstellen und ausführen

  1. Legen Sie die folgenden Umgebungsvariablen fest:

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

  2. 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)

  3. Führen Sie dieses Skript mit dem folgenden Befehl aus:

    (vm)$ python3 tpu-test.py

    Dieses Script führt eine Berechnung auf jedem TensorCore eines TPU-Pod-Slices 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),
  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.

  1. Trennen Sie die Verbindung zur Compute Engine:

    (vm)$ exit

  2. Löschen Sie Ihre Cloud TPU.

    $ gcloud compute tpus tpu-vm delete ${TPU_NAME} \
  --zone=${ZONE}

  3. Ü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=${ZONE}