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:

  1. TPU-Pod-Slice mit TensorFlow-Software erstellen
  2. Über SSH eine Verbindung zur TPU-VM herstellen
  3. 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

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

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