Menjalankan kode TensorFlow pada slice Pod TPU

Dokumen ini menunjukkan cara melakukan penghitungan sederhana menggunakan TensorFlow pada 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 skrip sederhana

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

Membuat slice Pod TPU v3-32 dengan runtime TensorFlow

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.16.1-pod-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.

Menghubungkan ke VM Cloud TPU menggunakan SSH

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

Membuat dan menjalankan skrip penghitungan sederhana

  1. Setel 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 tempel 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 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),
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 Anda.

  1. Putuskan koneksi dari Compute Engine:

    (vm)$ exit

  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 perintah berikut. Pastikan TPU Anda tidak lagi tercantum. Proses penghapusan mungkin memerlukan waktu beberapa menit.

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