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

Cloud TPU mit gcloud erstellen

1. Definieren Sie einige Umgebungsvariablen, um die Verwendung der Befehle zu erleichtern.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-east5-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite

   Beschreibungen von Umgebungsvariablen

   Variable	Beschreibung
   PROJECT_ID	Ihre Projekt-ID von Google Cloud . Verwenden Sie ein vorhandenes oder erstellen Sie ein neues Projekt.
   TPU_NAME	Der Name der TPU.
   ZONE	Die Zone, in der die TPU-VM erstellt werden soll. Weitere Informationen zu unterstützten Zonen finden Sie unter TPU-Regionen und -Zonen.
   ACCELERATOR_TYPE	Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
   RUNTIME_VERSION	Die Softwareversion der Cloud TPU.

2. Erstellen Sie Ihre TPU-VM mit dem folgenden Befehl:

   $ 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 keine SSH-Verbindung zu einer TPU-VM herstellen können, liegt das möglicherweise daran, dass die TPU-VM keine externe IP-Adresse hat. Wenn Sie ohne externe IP-Adresse auf eine TPU-VM 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] -f https://storage.googleapis.com/libtpu-releases/index.html

Prüfen, ob PyTorch auf TPUs zugreifen kann

Prüfen Sie mit dem folgenden Befehl, 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', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Einfache Berechnung durchführen

1. Erstellen Sie im aktuellen Verzeichnis eine Datei mit dem Namen tpu-test.py, kopieren Sie das folgende Script 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 Scripts 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')