Resnet50 auf Cloud TPU mit PyTorch trainieren

In dieser Anleitung erfahren Sie, wie Sie das ResNet-50-Modell auf einem Cloud TPU-Gerät mit PyTorch trainieren. Sie können dasselbe Muster auf andere TPU-optimierte Bildklassifikationsmodelle anwenden, die PyTorch und das ImageNet-Dataset verwenden.

Das Modell in dieser Anleitung basiert auf dem Framework Deep Residual Learning for Image Recognition, in dem erstmalig die Residualnetzwerkarchitektur (ResNet-Architektur) eingeführt wurde. In der Anleitung wird die 50-Layer-Variante ResNet-50 verwendet und das Training des Modells mit PyTorch/XLA veranschaulicht.

Ziele

  • Bereiten Sie das Dataset vor.
  • Trainingsjob ausführen
  • Ausgabeergebnisse überprüfen

Kosten

In dieser Anleitung werden kostenpflichtige Komponenten von Google Cloud verwendet, darunter:

  • Compute Engine
  • Cloud TPU

Sie können mithilfe des Preisrechners die Kosten für Ihre voraussichtliche Nutzung kalkulieren. Neuen Google Cloud-Nutzern steht möglicherweise eine kostenlose Testversion zur Verfügung.

Hinweise

Bevor Sie mit dieser Anleitung beginnen, prüfen Sie, ob Ihr Google Cloud-Projekt ordnungsgemäß eingerichtet ist.

  1. Melden Sie sich bei Ihrem Google-Konto an.

    Wenn Sie noch kein Konto haben, melden Sie sich hier für ein neues Konto an.

  2. Wählen Sie in der Cloud Console auf der Seite für die Projektauswahl ein Cloud-Projekt aus oder erstellen Sie eines.

    Zur Projektauswahl

  3. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein. So prüfen Sie, ob die Abrechnung für Ihr Projekt aktiviert ist.

  4. In dieser Anleitung werden kostenpflichtige Komponenten der Google Cloud verwendet. Rufen Sie die Seite mit den Cloud TPU-Preisen auf, um Ihre Kosten abzuschätzen. Denken Sie daran, nicht mehr benötigte Ressourcen zu bereinigen, um unnötige Kosten zu vermeiden.

Compute Engine-Instanz einrichten

  1. Öffnen Sie ein Cloud Shell-Fenster.

    Cloud Shell öffnen

  2. Erstellen Sie eine Variable für Ihre Projekt-ID.

    export PROJECT_ID=project-id
    
  3. Konfigurieren Sie das Befehlszeilentool gcloud, um das Projekt zu verwenden, in dem Sie Cloud TPU erstellen möchten.

    gcloud config set project ${PROJECT_ID}
    
  4. PyTorch-Version festlegen

    export VERSION=1.6
    
  5. Starten Sie die für diese Anleitung erforderliche Compute Engine-Ressource.

    gcloud compute instances create resnet50-tutorial \
    --zone=us-central1-a \
    --machine-type=n1-highmem-96  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=300GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform
    

    Wenn Sie Resnet50 mit echten Daten trainieren möchten, wählen Sie die Option Maschinentyp mit der höchsten Anzahl an CPUs aus. Resnet50 ist in der Regel sehr eingabegebunden, wodurch das Training ziemlich langsam sein kann. Dies kann vermieden werden, wenn viele Worker zum Einspeisen von Daten und ausreichend RAM zum Verwalten einer großen Anzahl von Worker-Threads vorhanden sind. Die besten Ergebnisse erhalten Sie mit dem Maschinentyp n1-highmem-96.

    Wenn Sie ImageNet herunterladen möchten, geben Sie eine Laufwerkgröße von mindestens 300 GB an. Wenn Sie nur fiktive Daten verwenden möchten, können Sie die standardmäßige Laufwerkgröße von 20 GB festlegen. In dieser Anleitung wird vorgeschlagen, beide Datasets zu verwenden.

Cloud TPU-Ressource starten

  1. Starten Sie über die virtuelle Compute Engine-Maschine mit dem folgenden Befehl eine Cloud TPU-Ressource:

    (vm) $ gcloud compute tpus create resnet50-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-${VERSION?}  \
    --accelerator-type=v3-8
    
  2. Ermitteln Sie die IP-Adresse für die Cloud TPU-Ressource.

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

PyTorch-Umgebung erstellen und konfigurieren

  1. Stellen Sie eine Verbindung zur neuen Compute Engine-Instanz her.

    gcloud compute ssh resnet50-tutorial --zone=us-central1-a
    

Ab jetzt bedeutet das Präfix (vm)$, dass Sie den Befehl in der Compute Engine-VM-Instanz ausführen sollten.

  1. Starten Sie eine conda-Umgebung.

    (vm) $ export VERSION=1.6
    (vm) $ conda activate torch-xla-${VERSION?}
    
  2. Konfigurieren Sie Umgebungsvariablen für die Cloud TPU-Ressource.

    (vm) $ export TPU_IP_ADDRESS=ip-address
    
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

Mit dem fiktiven Dataset trainieren

Wir empfehlen die Verwendung des fiktiven Datasets für die erste Ausführung anstelle der echten ImageNet-Daten, da fake_data automatisch auf der VM installiert wird und weniger Zeit und Ressourcen benötigt.

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=2 --batch_size=128 --log_steps=20

Mit dem echten Dataset trainieren

Wenn die Verwendung des Flags --fake_data erfolgreich ist, können Sie versuchen, mit echten Daten wie ImageNet zu trainieren.

test_train_mp_imagenet.py verwendet in der Regel torchvision.datasets.ImageFolder, sodass Sie jedes Dataset verwenden können, das richtig strukturiert ist. Weitere Informationen finden Sie in der Dokumentation zu ImageFolder.

Einige Vorschläge zu Befehlszeilenänderungen für die Verwendung realer Daten, vorausgesetzt, Sie haben das Dataset unter ~/imagenet gespeichert:

(vm) $ python /usr/share/torch-xla-${VERSION?}/pytorch/xla/test/test_train_mp_imagenet.py --datadir=~/imagenet --model=resnet50 --num_epochs=90 --num_workers=8 --batch_size=128 --log_steps=200

Zusätzliches CPU-Kontingent anfordern

Zusätzliche Ressourcen sollten einige Tage im Voraus geplant und angefordert werden, damit genügend Zeit für die Verarbeitung der Anfrage bleibt.

  1. Rufen Sie die Seite Kontingente auf.

    Zur Seite "Kontingente"

  2. Wählen Sie im Menü Dienst die Option Cloud TPU API aus.
  3. Wählen Sie die Region oder die Zonen aus, in denen Sie die CPUs verwenden möchten.
  4. Wählen Sie im Menü Messwert die Option "Keine" aus und geben Sie CPUs in das Suchfeld ein.
  5. Wählen Sie CPUs aus.
  6. Wählen Sie in der Liste Compute Engine API – CPUs aus und klicken Sie oben auf der Seite auf Kontingente bearbeiten.
  7. Geben Sie den gewünschten Kontingentbetrag und eine Beschreibung ein (erforderlich) und klicken Sie auf Fertig. Eine Anfrage wird zur Genehmigung an Ihren Internetanbieter gesendet.

Bereinigen

So vermeiden Sie, dass Ihrem Google Cloud Platform-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 user@projectname lauten und angeben, dass Sie sich in Cloud Shell befinden.

  2. Führen Sie in der Cloud Shell ctpu delete mit den Flags --zone und --name aus, die Sie beim Einrichten der Compute Engine-VM und der Cloud TPU verwendet haben. Dadurch werden sowohl Ihre VM als auch Ihre Cloud TPU gelöscht.

    $ ctpu delete --name=resnet50-tutorial \
      --zone=europe-west4-a
    
  3. Führen Sie ctpu status aus, um zu bestätigen, dass keine Instanzen zugeordnet wurden. So vermeiden Sie unnötige Gebühren für die TPU-Nutzung. Der Löschvorgang kann einige Minuten dauern. Eine Antwort wie die folgende gibt an, dass keine weiteren Instanzen vorhanden sind:

    $ ctpu status --zone=europe-west4-a
    
    2018/04/28 16:16:23 WARNING: Setting zone to "europe-west4-a"
    No instances currently exist.
            Compute Engine VM:     --
            Cloud TPU:             --
    

Nächste Schritte

Testen Sie die PyTorch Colabs: