Kurzanleitung zu Cloud TPU PyTorch

In dieser Kurzanleitung erfahren Sie, wie Sie eine Cloud TPU erstellen, PyTorch installieren und eine einfache Berechnung auf einer Cloud TPU ausführen. Eine ausführlichere Anleitung zum Trainieren eines Modells auf einer Cloud TPU finden Sie in einer der Cloud TPU-PyTorch-Anleitungen.

Hinweis

Bevor Sie mit dieser Kurzanleitung beginnen, müssen Sie ein Google Cloud Platform-Konto erstellen und das Google Cloud Platform SDK installieren. und konfigurieren Sie den Befehl gcloud. Weitere Informationen finden Sie unter Konto und Cloud TPU-Projekt einrichten.

Cloud TPU mit gcloud erstellen

Starten Sie mit dem Befehl gcloud eine Compute Engine-VM und eine Cloud TPU. Welchen Befehl Sie verwenden, hängt davon ab, ob Sie eine TPU-VM oder einen TPU-Knoten verwenden. Weitere Informationen zu den beiden VM-Architekturen finden Sie unter Systemarchitektur. Weitere Informationen zum Befehl gcloud finden Sie in der Referenz zu gcloud.

TPU-VM

$ gcloud alpha compute tpus tpu-vm create tpu-name \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=v2-alpha

Beschreibung der Befehls-Flags

zone
Die Zone, in der Sie die Cloud TPU erstellen möchten.
accelerator-type
Der Typ der zu erstellenden Cloud TPU.
version
Die Cloud TPU-Laufzeitversion.

TPU-Knoten

  1. Beim Erstellen eines TPU-Knotens für PyTorch erstellen Sie zuerst eine Compute Engine-VM-Instanz.

    gcloud compute instances create tpu-name \
    --zone=us-central1-b  \
    --machine-type=n1-standard-16  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=200GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform

    Beschreibung der Befehls-Flags

    project
    Ihre GCP-Projekt-ID
    name
    Der Name der zu erstellenden Cloud TPU.
    zone
    Die Zone, in der Sie die Cloud TPU erstellen möchten.
    disk-size
    Die Größe des Laufwerks in GB der VM, die mit dem Befehl gcloud erstellt wurde.
    machine-type
    Der Maschinentyp der zu erstellenden Compute Engine-VM.
    tf-version
    Die Version von TensorFlow gcloud wird auf der VM installiert.
    accelerator-type
    Der Typ der zu erstellenden Cloud TPU.
  2. Erstellen Sie als Nächstes die TPU-Instanz.

    gcloud compute tpus create tpu-name \
    --zone=us-central1-b \
    --network=default \
    --version=pytorch-1.9 \
    --accelerator-type=v3-8

    Beschreibung der Befehls-Flags

    name
    Der Name der zu erstellenden Cloud TPU.
    zone
    Die Zone, in der Sie die Cloud TPU erstellen möchten.
    network
    Das Netzwerk, mit dem die TPU verbunden ist.
    version
    Die Version von PyTorch gcloud wird auf der VM installiert.
    accelerator-type
    Der Typ der zu erstellenden Cloud TPU.

Verbindung zur Cloud TPU-VM herstellen

TPU-VM

$ gcloud alpha compute tpus tpu-vm ssh tpu-name \
  --zone us-central1-b

TPU-Knoten

gcloud compute ssh tpu-name --zone=us-central1-b

XRT-TPU-Gerätekonfiguration festlegen

TPU-VM

Konfigurieren Sie die Torch-XLA-Umgebung.

(vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

TPU-Knoten

  1. Ermitteln Sie die IP-Adresse des TPU-Knotens.

    (vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b
  2. Konfigurieren Sie die Torch-XLA-Umgebung. Ersetzen Sie your-tpu-ip-address durch die IP-Adresse Ihrer TPU.

    (vm)$ conda activate torch-xla-1.9
    (vm)$ export TPU_IP_ADDRESS=your-tpu-ip-address
    (vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

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 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')</span></pre>
    
    OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
    OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
    

Bereinigen

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

    TPU-VM

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

    TPU-Knoten

    $ gcloud compute tpus execution-groups delete tpu-name \
    --zone=us-central1-b

    Die Ausgabe dieses Befehls sollte bestätigen, dass Ihre TPU gelöscht wurde.

Nächste Schritte

Weitere Informationen zu Cloud TPU-VMs: