Menjalankan penghitungan pada VM Cloud TPU menggunakan PyTorch

Panduan memulai ini menunjukkan cara membuat Cloud TPU, menginstal PyTorch, dan menjalankan penghitungan sederhana di Cloud TPU. Untuk tutorial lebih mendalam yang menunjukkan cara melatih model di Cloud TPU, lihat salah satu Tutorial PyTorch Cloud TPU.

Sebelum memulai

Sebelum mengikuti panduan memulai ini, Anda harus membuat akun Google Cloud Platform, menginstal Google Cloud CLI, dan mengonfigurasi perintah gcloud. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan akun dan project Cloud TPU.

Membuat Cloud TPU dengan gcloud

Untuk membuat VM TPU dalam project pengguna, jaringan, dan zona default, jalankan:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-b \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

Deskripsi flag perintah

zone
Zona tempat Anda berencana membuat Cloud 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.
version
Versi software Cloud TPU .

Saat membuat TPU, Anda dapat meneruskan flag --network dan --subnetwork tambahan jika ingin menentukan jaringan dan subnetwork default. Jika tidak ingin menggunakan jaringan default, Anda harus meneruskan flag --network. Flag --subnetwork bersifat opsional dan dapat digunakan untuk menentukan subnetwork default untuk jaringan apa pun yang Anda gunakan (default atau ditentukan pengguna). Lihat gcloud halaman referensi API untuk mengetahui detail tentang tanda ini.

Hubungkan ke VM Cloud TPU Anda

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-b

Instal PyTorch/XLA di VM TPU Anda

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

Setel konfigurasi runtime TPU

Pastikan runtime PyTorch/XLA menggunakan TPU.

   (vm) $ export PJRT_DEVICE=TPU

Lakukan penghitungan sederhana:

  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)$ python3 tpu-test.py

    Output dari skrip menampilkan 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 dikenakan biaya untuk resource yang digunakan pada halaman ini, ikuti langkah-langkah berikut.

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

    (vm)$ exit

    Prompt 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 \
  --zone=us-central1-b

Output dari perintah ini akan mengonfirmasi bahwa TPU Anda telah dihapus.

Langkah selanjutnya

Baca lebih lanjut tentang VM Cloud TPU: