Melatih ResNet-50 di Cloud TPU dengan PyTorch

Tutorial ini menunjukkan cara melatih model ResNet-50 di perangkat Cloud TPU dengan PyTorch. Anda dapat menerapkan pola yang sama ke model klasifikasi gambar yang dioptimalkan untuk TPU lainnya yang menggunakan PyTorch dan set data ImageNet.

Model dalam tutorial ini didasarkan pada Deep Residual Learning for Image Recognition, yang pertama kali memperkenalkan arsitektur jaringan residual (ResNet). Tutorial ini menggunakan varian 50 lapisan, ResNet-50, dan menunjukkan cara melatih model menggunakan PyTorch/XLA.

Tujuan

  • Siapkan set data.
  • Jalankan tugas pelatihan.
  • Verifikasi hasil output.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponen Google Cloudyang dapat ditagih berikut:

  • Compute Engine
  • Cloud TPU

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Pengguna Google Cloud baru mungkin memenuhi syarat untuk mendapatkan uji coba gratis.

Sebelum memulai

Sebelum memulai tutorial ini, periksa apakah project Google Cloud Anda telah disiapkan dengan benar.

  1. Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.
  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  3. Verify 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.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  5. Verify that billing is enabled for your Google Cloud project.

  6. Panduan ini menggunakan komponen Google Cloudyang dapat ditagih. Lihat halaman harga Cloud TPU untuk memperkirakan biaya Anda. Pastikan untuk membersihkan resource yang Anda buat setelah Anda selesai menggunakannya untuk menghindari tagihan yang tidak perlu.

Membuat VM TPU

  1. Buka jendela Cloud Shell atau terminal.

    Buka Cloud Shell

  2. Tentukan variabel lingkungan untuk parameter perintah Google Cloud CLI:

    export PROJECT_ID=your-project-id
    export TPU_NAME=your-tpu-name
    export ZONE=us-central1-b
    export ACCELERATOR_TYPE=v6e-8
    export RUNTIME_VERSION=v2-alpha-tpuv6e

    Deskripsi variabel lingkungan

    • PROJECT_ID: Project ID Google Cloud Anda. Gunakan project yang sudah ada atau buat project baru.
    • TPU_NAME: Nama TPU.
    • ZONE: Zona tempat membuat VM TPU. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
    • ACCELERATOR_TYPE: Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat Versi TPU.
    • RUNTIME_VERSION: Versi software Cloud TPU.

  3. Membuat VM TPU

     gcloud compute tpus tpu-vm create $TPU_NAME \
     --accelerator-type=$ACCELERATOR_TYPE \
     --version=$RUNTIME_VERSION \
     --zone=$ZONE \
     --project=$PROJECT_ID
    
    Google Cloud
  4. Hubungkan ke VM TPU Anda menggunakan SSH:

     gcloud compute tpus tpu-vm ssh  $TPU_NAME --zone=$ZONE
    

Menginstal PyTorch/XLA di VM TPU

  1. Buat file bernama requirements.txt dan tambahkan konten berikut:

    --find-links https://storage.googleapis.com/libtpu-releases/index.html
    --find-links https://storage.googleapis.com/libtpu-wheels/index.html
    torch~=2.6.0
    torch_xla[tpu]~=2.6.0
    torchvision
    ray[default]==2.40.0
    
  2. Instal dependensinya:

    pip install -r requirements.txt
    
  3. Clone repositori GitHub PyTorch/XLA

    git clone --depth=1 https://github.com/pytorch/xla.git
    
  4. Jalankan skrip pelatihan dengan data palsu

    PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py \
    --fake_data \
    --batch_size=256 \
    --num_epochs=1
    

Pembersihan

Agar tidak perlu membayar biaya pada akun Google Cloud Anda untuk resource yang digunakan dalam tutorial ini, hapus project yang berisi resource tersebut, atau simpan project dan hapus setiap resource.

  1. Putuskan koneksi dari VM TPU:

    exit
    
  2. Hapus VM TPU Anda.

    gcloud compute tpus tpu-vm delete $TPU_NAME \
        --zone=$ZONE
    

Langkah berikutnya