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.

Peran yang diperlukan

Agar mendapatkan izin yang diperlukan untuk membuat TPU dan terhubung ke TPU menggunakan SSH, minta administrator untuk memberi Anda peran IAM berikut di project Anda:

Untuk mengetahui informasi selengkapnya tentang pemberian peran, lihat Mengelola akses ke project, folder, dan organisasi.

Anda mungkin juga bisa mendapatkan izin yang diperlukan melalui peran khusus atau peran bawaan lainnya.

Buat Cloud TPU menggunakan gcloud

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

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

    Deskripsi variabel lingkungan

    Variabel Deskripsi
    PROJECT_ID ID project Google Cloud Anda. Gunakan project yang ada atau buat project baru.
    TPU_NAME Nama TPU.
    ZONE Zona tempat VM TPU akan dibuat. 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.

  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

Jika Anda gagal terhubung ke VM TPU menggunakan SSH, hal ini mungkin karena VM TPU tidak memiliki alamat IP eksternal. Untuk mengakses VM TPU tanpa alamat IP eksternal, ikuti petunjuk di Menghubungkan ke VM TPU tanpa alamat IP publik.

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] -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', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Melakukan penghitungan dasar

  1. Buat file bernama tpu-test.py di direktori saat ini, lalu salin dan tempel 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 pada halaman ini, ikuti langkah-langkah berikut.

  1. Putuskan koneksi dari instance Cloud TPU, jika Anda belum melakukannya:

    (vm)$ exit

    Sekarang perintah Anda akan menjadi username@projectname, yang menunjukkan Anda berada dalam Cloud Shell.

  2. Hapus Cloud TPU Anda.

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

  3. Pastikan resource telah dihapus dengan menjalankan perintah berikut. Pastikan TPU Anda tidak lagi tercantum. Proses penghapusan mungkin memerlukan waktu beberapa menit.

    $ gcloud compute tpus tpu-vm list \
    --zone=$ZONE

Langkah berikutnya

Baca selengkapnya tentang VM Cloud TPU: