Menjalankan penghitungan di VM Cloud TPU menggunakan PyTorch

Dokumen ini memberikan pengantar singkat tentang cara menggunakan PyTorch dan Cloud TPU.

Sebelum memulai

Sebelum menjalankan perintah dalam dokumen ini, Anda harus membuat akun Google Cloud , menginstal Google Cloud CLI, dan mengonfigurasi perintah gcloud. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan lingkungan Cloud TPU.

Membuat Cloud TPU menggunakan gcloud

  1. Tentukan beberapa variabel lingkungan agar perintah lebih mudah digunakan.

    export PROJECT_ID=your-project
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name

    Deskripsi variabel lingkungan

    PROJECT_ID
    Project ID Google Cloud Anda.
    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.
    ZONE
    Zona tempat Anda berencana membuat Cloud TPU.
    RUNTIME_VERSION
    Versi runtime Cloud TPU.
    TPU_NAME
    Nama yang ditetapkan pengguna untuk Cloud TPU Anda.

  2. Buat VM TPU dengan menjalankan perintah berikut:

    $ gcloud compute tpus tpu-vm create $TPU_NAME \
--project=$PROJECT_ID \
--zone=$ZONE \
--accelerator-type=$ACCELERATOR_TYPE \
--version=$RUNTIME_VERSION

Menghubungkan ke VM Cloud TPU

Hubungkan ke VM TPU Anda melalui SSH menggunakan perintah berikut:

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

Menginstal PyTorch/XLA di VM TPU

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html

Memverifikasi bahwa PyTorch dapat mengakses TPU

Gunakan perintah berikut untuk memverifikasi bahwa PyTorch dapat mengakses TPU Anda.

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

Output dari perintah akan terlihat seperti berikut:

['xla:0', 'xla:1', 'xla:2', 'xla:3']

Melakukan penghitungan dasar

  1. Buat file bernama tpu-test.py di direktori saat ini, lalu salin dan tempelkan skrip berikut ke dalamnya.

    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. Jalankan skrip:

    (vm)$ PJRT_DEVICE=TPU python3 tpu-test.py

    Output dari skrip menunjukkan hasil komputasi:

    tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

Pembersihan

Agar akun Google Cloud Anda tidak dikenai biaya untuk resource yang digunakan di halaman ini, ikuti langkah-langkah berikut.

  1. Putuskan koneksi dari instance Compute Engine, jika Anda belum melakukannya:

    (vm)$ exit

    Perintah Anda sekarang akan menjadi username@projectname, yang menunjukkan bahwa Anda berada di Cloud Shell.

  2. Hapus Cloud TPU Anda.

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

Output perintah ini akan mengonfirmasi bahwa TPU Anda telah dihapus.

Langkah berikutnya

Baca selengkapnya tentang VM Cloud TPU: