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 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. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

  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 zu Ihrer TPU-VM her:

    gcloud compute tpus tpu-vm ssh  your-tpu-name --zone=us-central2-b
  4. So installieren Sie PyTorch/XLA auf Ihrer TPU-VM:

    (vm)$ pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html
  5. PyTorch/XLA-GitHub-Repository klonen

    (vm)$ git clone --depth=1 --branch r2.5 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 Befehl für das Trainingsskript gibt das Flag --datadir den Speicherort des Datasets an, das trainiert werden soll. Im folgenden Befehl wird davon ausgegangen, dass sich der 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 Ihre TPU-VM.

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

Nächste Schritte