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

In diesem Dokument erfahren Sie, wie Sie mit TensorFlow eine Berechnung 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 ein Dienstkonto nach Berechtigungen zum Aufrufen der Cloud TPU API. Standardmäßig verwendet Ihre TPU-VM das Standarddienstkonto der Compute Engine, das alle erforderlichen Cloud TPU-Berechtigungen enthält. Wenn Sie Ihre eigene Dienstkonto, dem Sie den TPU Viewer hinzufügen müssen Rolle zu Ihrem Dienstkonto. Weitere Informationen zu Google Cloud-Rollen finden Sie unter Informationen zu Rollen. Sie können Ihr eigenes Dienstkonto mit dem Flag --service-account angeben, wenn die TPU-VM erstellen.

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.17.0-pod-pjrt
    $ export ACCELERATOR_TYPE=v3-32

TPU v3-32-Pod-Slice 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
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 jede TPU-Version finden Sie unter TPU-Versionen.
version
Die Softwareversion von Cloud TPU.

Ü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 Skript führt eine Berechnung für jeden TensorCore eines TPU-Pod-Slice 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}