Mit PyTorch eine Berechnung auf einer Cloud TPU-VM ausführen

In dieser Kurzanleitung wird gezeigt, wie Sie eine Cloud TPU erstellen, PyTorch installieren und eine einfache Berechnung auf einer Cloud TPU ausführen. Ein ausführlicheres Tutorial, in dem Sie wie Sie ein Modell auf einer Cloud TPU trainieren, Anleitungen zu Cloud TPU PyTorch

Hinweise

Bevor Sie dieser Kurzanleitung folgen, müssen Sie ein Google Cloud-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 mit gcloud erstellen

So erstellen Sie eine TPU-VM im Standardprojekt, -Netzwerk und -Zone des Nutzers:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

Beschreibung der Befehls-Flags

zone
Die Zone in dem Sie Ihre 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 Version der Cloud TPU-Software.

Beim Erstellen Ihrer TPU können Sie die zusätzlichen Flags --network und --subnetwork übergeben, wenn Sie das Standardnetzwerk und das Standardsubnetz angeben möchten. Wenn Sie das Standardnetzwerk nicht verwenden möchten, müssen Sie den Flag --network. Das Flag --subnetwork ist optional und kann verwendet werden, um ein Standard-Subnetzwerk für das verwendete Netzwerk (Standard- oder vom Nutzer angegeben) anzugeben. Weitere Informationen zu diesen Flags findest du auf der gcloudAPI-Referenzseite.

Verbindung zur Cloud TPU-VM herstellen

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

PyTorch/XLA auf der TPU-VM installieren

   (vm)$ pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

TPU-Laufzeitkonfiguration festlegen

Achten Sie darauf, dass die PyTorch/XLA-Laufzeit die TPU verwendet.

   (vm) $ export PJRT_DEVICE=TPU

Führen Sie eine einfache Berechnung durch:

  1. 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 torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

  2. Führen Sie das Skript aus:

      (vm)$ python3 tpu-test.py

    Die Ausgabe des Skripts zeigt das Ergebnis der Berechnung:

    tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

Bereinigen

Mit den folgenden Schritten vermeiden Sie, dass Ihrem Google Cloud-Konto die in dieser Anleitung 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.

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central1-a

Die Ausgabe dieses Befehls sollte ergeben, dass Ihre TPU gelöscht wurde.

Nächste Schritte

Weitere Informationen zu Cloud TPU-VMs: