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.

Lernziele

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

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

  • Compute Engine
  • Cloud TPU

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen. 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 Cloud-Konto an. Wenn Sie mit Google Cloud noch nicht vertraut sind, erstellen Sie ein Konto, um die Leistungsfähigkeit unserer Produkte in der Praxis sehen und bewerten zu können. Neukunden erhalten außerdem ein Guthaben von 300 $, um Arbeitslasten auszuführen, zu testen und bereitzustellen.

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

    Zur Projektauswahl

  3. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

  4. Wählen Sie in der Google Cloud Console auf der Seite der Projektauswahl ein Google Cloud-Projekt aus oder erstellen Sie eines.

    Zur Projektauswahl

  5. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

    6. 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.

TPU-VM erstellen

  1. Öffnen Sie ein Cloud Shell-Fenster.

    Cloud Shell öffnen

  2. TPU-VM erstellen 

    gcloud compute tpus tpu-vm create your-tpu-name \
--accelerator-type=v4-8 \
--version=tpu-ubuntu2204-base \
--zone=us-central2-b \
--project=your-project

  3. Stellen Sie über SSH eine Verbindung zur TPU-VM her:

    gcloud compute tpus tpu-vm ssh  your-tpu-name --zone=us-central2-b

  4. Installieren Sie PyTorch/XLA auf Ihrer TPU-VM:

    (vm)$ pip install torch~=2.1.0 torch_xla[tpu]~=2.1.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

  5. Klonen Sie das PyTorch/XLA-GitHub-Repository.

    (vm)$ git clone --depth=1 --branch r2.1 https://github.com/pytorch/xla.git

  6. Trainingsskript mit fiktiven Daten ausführen

    (vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1

Wenn Sie das Modell mit fiktiven Daten trainieren können, können Sie versuchen, mit echten Daten wie ImageNet zu trainieren. Eine Anleitung zum Herunterladen von ImageNet finden Sie unter ImageNet herunterladen. Im Trainingsskriptbefehl gibt das Flag --datadir den Speicherort des Datasets an, auf dem das Training ausgeführt werden soll. Bei dem folgenden Befehl wird davon ausgegangen, dass sich das ImageNet-Dataset in ~/imagenet befindet.

   (vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py  --datadir=~/imagenet --batch_size=256 --num_epochs=1

Bereinigen

Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.

  1. Trennen Sie die Verbindung zur TPU-VM:

    (vm) $ exit

    Die Eingabeaufforderung sollte nun username@projectname lauten und angeben, dass Sie sich in Cloud Shell befinden.

  2. Löschen Sie die TPU-VM.

    $ gcloud compute tpus tpu-vm delete resnet50-tutorial \
   --zone=us-central2-b

Nächste Schritte

Testen Sie die PyTorch Colabs: