Menjalankan penghitungan di VM Cloud TPU menggunakan TensorFlow

Panduan memulai ini menunjukkan cara membuat Cloud TPU, menginstal TensorFlow, dan menjalankan penghitungan sederhana di Cloud TPU. Untuk tutorial yang lebih mendalam yang menunjukkan cara melatih model di Cloud TPU, lihat salah satu Tutorial 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 VM Cloud TPU atau Node dengan gcloud

Luncurkan Cloud TPU Compute Engine menggunakan perintah gcloud. Perintah yang digunakan bergantung pada apakah Anda menggunakan VM TPU atau node TPU. Untuk informasi selengkapnya tentang dua arsitektur VM, lihat Arsitektur Sistem. Untuk mengetahui informasi selengkapnya tentang perintah gcloud, lihat referensi gcloud.

VM TPU

$ gcloud compute tpus tpu-vm create tpu-name \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.16.1-pjrt

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.

Node TPU

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.12.0 \

Deskripsi flag perintah

project
ID project Google Cloud Anda
name
Nama Cloud TPU yang akan dibuat.
zone
Zona tempat Anda berencana membuat Cloud TPU.
disk-size
Ukuran hard disk dalam GB VM yang dibuat oleh perintah gcloud.
machine-type
Jenis mesin VM Compute Engine yang akan dibuat.
tf-version
Versi TensorFlow gcloud diinstal di VM. Lihat
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.

Hubungkan ke VM Cloud TPU Anda

Saat menggunakan VM TPU, Anda harus secara eksplisit terhubung ke VM TPU menggunakan SSH. Saat menggunakan TPU Node, Anda harus otomatis terhubung ke SSH ke VM Compute Engine. Jika Anda tidak terhubung secara otomatis, gunakan perintah berikut.

VM TPU

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone europe-west4-a

Node TPU

$ gcloud compute ssh tpu-name \
    --zone=europe-west4-a

Menjalankan contoh sederhana menggunakan TensorFlow

VM TPU

Setelah Anda terhubung ke VM TPU, tetapkan variabel lingkungan berikut.

  (vm)$ export TPU_NAME=local

Saat membuat TPU, jika Anda menetapkan parameter --version ke versi yang diakhiri dengan -pjrt, tetapkan variabel lingkungan berikut untuk mengaktifkan runtime PJRT: 

  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

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

  import tensorflow as tf
  print("Tensorflow version " + tf.__version__)

  @tf.function
  def add_fn(x,y):
    z = x + y
    return z

  cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
  tf.config.experimental_connect_to_cluster(cluster_resolver)
  tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
  strategy = tf.distribute.TPUStrategy(cluster_resolver)

  x = tf.constant(1.)
  y = tf.constant(1.)
  z = strategy.run(add_fn, args=(x,y))
  print(z)

Node TPU

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

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)

@tf.function
def add_fn(x,y):
    z = x + y
    return z

x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)

Jalankan skrip ini dengan perintah berikut:

(vm)$ python3 tpu-test.py

Skrip ini menjalankan komputasi sederhana pada setiap TensorCore TPU. Output-nya akan terlihat seperti berikut:

PerReplica:{
  0: tf.Tensor(2.0, shape=(), dtype=float32),
  1: tf.Tensor(2.0, shape=(), dtype=float32),
  2: tf.Tensor(2.0, shape=(), dtype=float32),
  3: tf.Tensor(2.0, shape=(), dtype=float32),
  4: tf.Tensor(2.0, shape=(), dtype=float32),
  5: tf.Tensor(2.0, shape=(), dtype=float32),
  6: tf.Tensor(2.0, shape=(), dtype=float32),
  7: tf.Tensor(2.0, shape=(), dtype=float32)
}

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.

    VM TPU

    $ gcloud compute tpus tpu-vm delete tpu-name \
--zone=europe-west4-a

    Node TPU

    $ gcloud compute tpus execution-groups delete tpu-name \
--zone=europe-west4-a

  3. Verifikasi bahwa resource telah dihapus dengan menjalankan gcloud compute tpus tpu-vm list. Penghapusan memerlukan waktu beberapa menit.

    VM TPU

    $ gcloud compute tpus tpu-vm list --zone=europe-west4-a

    Node TPU

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a

Langkah selanjutnya

Untuk mengetahui informasi selengkapnya tentang Cloud TPU, lihat: