Berechnung mit PyTorch auf einer Cloud TPU-VM ausführen

In diesem Dokument erhalten Sie eine kurze Einführung in die Arbeit mit PyTorch und Cloud TPU.

Hinweise

Bevor Sie die Befehle in diesem Dokument ausführen, müssen Sie ein Google Cloud -Konto erstellen, die Google Cloud CLI installieren und den gcloud-Befehl konfigurieren. Weitere Informationen finden Sie unter Cloud TPU-Umgebung einrichten.

Cloud TPU mit gcloud erstellen

  1. Definieren Sie einige Umgebungsvariablen, um die Befehle nutzerfreundlicher zu gestalten.

    export PROJECT_ID=your-project
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name
    PROJECT_ID
    Ihre Google Cloud Projekt-ID.
    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.
    ZONE
    Die Zone, in der Sie die Cloud TPU erstellen möchten.
    RUNTIME_VERSION
    Die Version der Cloud TPU-Laufzeit.
    TPU_NAME
    Der vom Nutzer zugewiesene Name für Ihre Cloud TPU.

  2. Führen Sie den folgenden Befehl aus, um die TPU-VM zu erstellen:

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

Verbindung zur Cloud TPU-VM herstellen

Stellen Sie mit dem folgenden Befehl eine SSH-Verbindung zu Ihrer TPU-VM her:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Wenn Sie über SSH keine Verbindung zu einer TPU-VM herstellen können, hat die TPU-VM möglicherweise keine externe IP-Adresse. Wenn Sie auf eine TPU-VM ohne externe IP-Adresse zugreifen möchten, folgen Sie der Anleitung unter Verbindung zu einer TPU-VM ohne öffentliche IP-Adresse herstellen.

PyTorch/XLA auf der TPU-VM installieren

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html

Prüfen, ob PyTorch auf TPUs zugreifen kann

Mit dem folgenden Befehl können Sie prüfen, ob PyTorch auf Ihre TPUs zugreifen kann.

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

Die Ausgabe des Befehls sollte so aussehen:

['xla:0', 'xla:1', 'xla:2', 'xla:3']

Einfache Berechnungen durchführen

  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)$ PJRT_DEVICE=TPU 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 Cloud TPU-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 \
    --project=$PROJECT_ID \
    --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

Nächste Schritte

Weitere Informationen zu Cloud TPU-VMs: