Menjalankan kode TensorFlow pada slice Pod TPU

Dokumen ini menunjukkan cara melakukan penghitungan menggunakan TensorFlow di Pod TPU. Anda akan melakukan langkah-langkah berikut:

  1. Membuat slice Pod TPU dengan software TensorFlow
  2. Menghubungkan ke VM TPU menggunakan SSH
  3. Membuat dan menjalankan contoh skrip

VM TPU mengandalkan Akun Layanan untuk mendapatkan izin guna memanggil Cloud TPU API. Secara default, VM TPU Anda akan menggunakan akun layanan Compute Engine default yang menyertakan semua izin Cloud TPU yang diperlukan. Jika menggunakan akun layanan Anda sendiri, Anda harus menambahkan peran TPU Viewer ke akun layanan Anda. Untuk mengetahui informasi selengkapnya tentang peran Google Cloud, lihat Memahami peran. Anda dapat menentukan akun layanan Anda sendiri menggunakan flag --service-account saat membuat VM TPU.

Menyiapkan lingkungan Anda

  1. Di Cloud Shell, jalankan perintah berikut untuk memastikan Anda menjalankan gcloud versi saat ini:

    $ gcloud components update

    Jika Anda perlu menginstal gcloud, gunakan perintah berikut:

    $ sudo apt install -y google-cloud-sdk

  2. Buat beberapa variabel lingkungan:

    $ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=europe-west4-a
$ export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

Membuat slice Pod TPU v3-32 dengan runtime TensorFlow

$ gcloud compute tpus tpu-vm create ${TPU_NAME}} \
  --zone=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE} \
  --version=${RUNTIME_VERSION}

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.

Menghubungkan ke VM Cloud TPU menggunakan SSH

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --zone=${ZONE}

Membuat dan menjalankan contoh skrip

  1. Tetapkan variabel lingkungan berikut.

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

  2. 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__)

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

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

@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)

  3. Jalankan skrip ini dengan perintah berikut:

    (vm)$ python3 tpu-test.py

    Skrip ini melakukan penghitungan pada setiap TensorCore dari slice Pod 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),
  8: tf.Tensor(2.0, shape=(), dtype=float32),
  9: tf.Tensor(2.0, shape=(), dtype=float32),
  10: tf.Tensor(2.0, shape=(), dtype=float32),
  11: tf.Tensor(2.0, shape=(), dtype=float32),
  12: tf.Tensor(2.0, shape=(), dtype=float32),
  13: tf.Tensor(2.0, shape=(), dtype=float32),
  14: tf.Tensor(2.0, shape=(), dtype=float32),
  15: tf.Tensor(2.0, shape=(), dtype=float32),
  16: tf.Tensor(2.0, shape=(), dtype=float32),
  17: tf.Tensor(2.0, shape=(), dtype=float32),
  18: tf.Tensor(2.0, shape=(), dtype=float32),
  19: tf.Tensor(2.0, shape=(), dtype=float32),
  20: tf.Tensor(2.0, shape=(), dtype=float32),
  21: tf.Tensor(2.0, shape=(), dtype=float32),
  22: tf.Tensor(2.0, shape=(), dtype=float32),
  23: tf.Tensor(2.0, shape=(), dtype=float32),
  24: tf.Tensor(2.0, shape=(), dtype=float32),
  25: tf.Tensor(2.0, shape=(), dtype=float32),
  26: tf.Tensor(2.0, shape=(), dtype=float32),
  27: tf.Tensor(2.0, shape=(), dtype=float32),
  28: tf.Tensor(2.0, shape=(), dtype=float32),
  29: tf.Tensor(2.0, shape=(), dtype=float32),
  30: tf.Tensor(2.0, shape=(), dtype=float32),
  31: tf.Tensor(2.0, shape=(), dtype=float32)
}

Pembersihan

Setelah selesai menggunakan VM TPU, ikuti langkah-langkah berikut untuk membersihkan resource.

  1. Putuskan koneksi dari Compute Engine:

    (vm)$ exit

  2. Hapus Cloud TPU Anda.

    $ gcloud compute tpus tpu-vm delete ${TPU_NAME} \
  --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}