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
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.
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
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)
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.
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.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: