TensorFlow를 사용하여 Cloud TPU VM에서 계산 실행

이 빠른 시작에서는 Cloud TPU를 만들고 TensorFlow를 사용하여 Cloud TPU에서 계산을 실행하는 방법을 보여줍니다. Cloud TPU에서 모델을 학습시키는 방법을 보여주는 자세한 튜토리얼은 Cloud TPU 튜토리얼 중 하나를 참고하세요.

시작하기 전에

이 빠른 시작을 수행하기 전에 Google Cloud 계정을 만들고 Google Cloud CLI를 설치하고 gcloud 명령어를 구성해야 합니다. 자세한 내용은 계정 및 Cloud TPU 프로젝트 설정을 참조하세요.

gcloud를 사용하여 Cloud TPU VM 만들기

gcloud 명령어를 사용하여 Cloud TPU를 생성합니다.

$ 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

명령어 플래그 설명

tpu-name
생성할 Cloud TPU의 이름입니다.
zone
Cloud TPU를 만들려는 영역입니다.
accelerator-type
가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
version
TPU 런타임 버전입니다. 이 빠른 시작에 사용되는 버전에는 TensorFlow가 사전 설치되어 있습니다.
project
Cloud TPU를 만들려는 Google Cloud CLI 프로젝트의 이름입니다.

gcloud 명령어에 대한 자세한 내용은 gcloud 참조를 확인하세요.

Cloud TPU VM에 연결

SSH를 사용하여 TPU VM에 연결합니다.

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

TensorFlow가 TPU에 액세스할 수 있는지 확인

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

    import tensorflow as tf

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

  2. 스크립트를 실행합니다.

    (vm)$ python3 tpu-count.py

    스크립트의 출력에는 TPU VM에서 사용할 수 있는 TPU 코어 수가 표시됩니다.

    TensorFlow can access 8 TPU cores

TensorFlow를 사용하여 기본 연산 실행

TPU VM에 연결되면 다음 환경 변수를 설정합니다.

  (vm)$ export TPU_NAME=local

TPU를 만들 때 --version 매개변수를 -pjrt로 끝나는 버전으로 설정한 경우 다음 환경 변수를 설정하여 PJRT 런타임을 사용 설정합니다. 

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

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

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)

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

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

삭제

이 페이지에서 사용한 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 다음 단계를 수행합니다.

  1. Compute Engine 인스턴스에서 연결을 해제합니다.

    (vm)$ exit

    프롬프트가 username@projectname으로 바뀌면 Cloud Shell에 있는 것입니다.

  2. Cloud TPU를 삭제합니다.

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

  3. gcloud compute tpus tpu-vm list를 실행하여 리소스가 삭제되었는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다.

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

다음 단계

Cloud TPU에 대한 자세한 내용은 다음을 참조하세요.