TPU Pod 슬라이스에서 TensorFlow 코드 실행

TensorFlow를 실행하는 TPU VM Pod 설정 및 계산 실행

  1. 다음 절차에 따라 TensorFlow를 실행하는 포드를 설정하고 계산을 실행합니다.

    $ gcloud compute tpus tpu-vm create tpu-name \
--zone=europe-west4-a \
--accelerator-type=v3-32 \
--version=tpu-vm-tf-2.11.0-pod

    명령어 플래그 설명

    zone
    Cloud TPU를 만들려는 영역입니다.
    accelerator-type
    생성할 Cloud TPU의 유형입니다.
    version
    Cloud TPU 소프트웨어 버전입니다. 기본값은 최신 TensorFlow 소프트웨어 버전입니다.

  2. Cloud TPU VM에 연결

    Cloud TPU Pod의 TPU VM에 SSH를 통해 연결합니다.

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

  3. 다음 환경 변수를 설정합니다.

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

  4. 현재 디렉터리에 tpu-test.py이라는 파일을 만들고 이 파일에 다음 스크립트를 복사하여 붙여넣습니다.

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)

다음 명령어를 사용하여 이 스크립트를 실행합니다.

(vm)$ python3 tpu-test.py

이 스크립트는 TPU의 각 TensorCore에서 간단한 계산을 수행합니다. 출력은 다음과 비슷합니다. 

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

삭제

TPU VM 사용이 완료되었으면 다음 단계에 따라 리소스를 삭제하세요.

  1. Compute Engine의 연결을 해제합니다.

    (vm)$ exit

  2. Cloud TPU를 삭제합니다.

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

  3. 다음 명령어를 실행하여 리소스가 삭제되었는지 확인합니다. TPU가 더 이상 나열되지 않았는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다.

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