Menjalankan penghitungan di VM Cloud TPU menggunakan TensorFlow

Panduan memulai ini menunjukkan cara membuat Cloud TPU dan menjalankan penghitungan di Cloud TPU menggunakan TensorFlow. 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, menginstal Google Cloud CLI, dan mengonfigurasi perintah gcloud. Untuk mengetahui informasi selengkapnya, lihat Menyiapkan akun dan project Cloud TPU.

Membuat VM Cloud TPU dengan gcloud

Buat Cloud TPU menggunakan perintah gcloud.

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

Deskripsi flag perintah

tpu-name
Nama Cloud TPU yang akan dibuat.
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 runtime TPU. Versi yang digunakan dalam panduan memulai ini telah menginstal TensorFlow sebelumnya.
project
Nama project Google Cloud CLI tempat Anda membuat Cloud TPU.

Untuk mengetahui informasi selengkapnya tentang perintah gcloud, lihat referensi gcloud.

Menghubungkan ke VM Cloud TPU

Hubungkan ke VM TPU Anda menggunakan SSH:

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone europe-west4-a \
  --project=your-gcp-project-name

Memverifikasi bahwa TensorFlow dapat mengakses TPU

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

    import tensorflow as tf

print(f"TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores")

  2. Jalankan skrip:

    (vm)$ python3 tpu-count.py

    Output dari skrip menunjukkan jumlah core TPU yang tersedia untuk VM TPU:

    TensorFlow can access 8 TPU cores

Menjalankan komputasi dasar menggunakan TensorFlow

Setelah 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 tempelkan 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)

Jalankan skrip ini dengan perintah berikut:

(vm)$ python3 tpu-test.py

Skrip ini melakukan komputasi pada setiap TensorCore TPU. Outputnya 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 tidak menimbulkan biaya pada akun Google Cloud Anda untuk resource yang digunakan pada 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 \
--zone=europe-west4-a

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

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

Langkah selanjutnya

Untuk mengetahui informasi selengkapnya tentang Cloud TPU, lihat: