Berechnung mit PyTorch 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. Eine ausführliche Anleitung zum Trainieren eines Modells auf einer Cloud TPU finden Sie in einer der PyTorch-Anleitungen für Cloud TPU.
Hinweise
Bevor Sie dieser Kurzanleitung folgen, müssen Sie ein Google Cloud Platform-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 in der Ausführung des Standardprojekts, -netzwerks und der -zone:
$ gcloud compute tpus tpu-vm create tpu-name \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=tpu-ubuntu2204-base
Beschreibung der Befehls-Flags
zone
- Die Zone, in der Sie Ihre 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 die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
version
- Die Version der Cloud TPU-Software.
Beim Erstellen der TPU können Sie die zusätzlichen Flags --network
und --subnetwork
übergeben, wenn Sie das Standardnetzwerk und -subnetzwerk angeben möchten.
Wenn Sie das Standardnetzwerk nicht verwenden möchten, müssen Sie das Flag --network
übergeben. Das Flag --subnetwork
ist optional und kann verwendet werden, um ein Standard-Subnetzwerk für das von Ihnen verwendete Netzwerk anzugeben (Standard oder benutzerdefiniert). Weitere Informationen zu diesen Flags finden Sie auf der
API-Referenzseite gcloud
.
Verbindung zur Cloud TPU-VM herstellen
$ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-b
PyTorch/XLA auf der TPU-VM installieren
(vm)$ pip install torch~=2.2.0 torch_xla[tpu]~=2.2.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:
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)
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 auf dieser Seite verwendeten Ressourcen in Rechnung gestellt werden.
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.Löschen Sie Ihre Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone=us-central1-b
Die Ausgabe dieses Befehls sollte ergeben, dass Ihre TPU gelöscht wurde.
Nächste Schritte
Weitere Informationen zu Cloud TPU-VMs: