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

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

시작하기 전에

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

gcloud로 Cloud TPU VM 또는 노드 만들기

gcloud 명령어를 사용하여 Compute Engine Cloud TPU를 실행합니다. 사용하는 명령어는 TPU VM과 TPU 노드 중 무엇을 사용하는지에 따라 다릅니다. 두 가지 VM 아키텍처에 관한 자세한 내용은 시스템 아키텍처를 참조하세요. gcloud 명령어에 대한 자세한 내용은 gcloud 참조를 확인하세요.

TPU VM

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

명령어 플래그 설명

zone
Cloud TPU를 만들려는 영역입니다.
accelerator-type
가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
version
Cloud TPU 소프트웨어 버전입니다.

TPU 노드

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.12.0 \

명령어 플래그 설명

project
: Google Cloud 프로젝트 ID입니다.
name
생성할 Cloud TPU의 이름입니다.
zone
Cloud TPU를 만들려는 영역입니다.
disk-size
gcloud 명령어로 만든 VM의 하드 디스크 크기(GB)입니다.
machine-type
생성할 Compute Engine VM의 머신 유형입니다.
tf-version
TensorFlow gcloud 버전이 VM에 설치됩니다. 다음을 참조하세요.
accelerator-type
가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.

Cloud TPU VM에 연결

TPU VM을 사용할 때 SSH를 사용하여 TPU VM에 명시적으로 연결해야 합니다. TPU 노드를 사용하면 Compute Engine VM에 SSH를 통해 자동 연결됩니다. 자동으로 연결되지 않으면 다음 명령어를 사용합니다.

TPU VM

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

TPU 노드

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

TensorFlow를 사용하여 간단한 예시 실행

TPU VM

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)

TPU 노드

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

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

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

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

삭제

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

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

    (vm)$ exit

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

  2. Cloud TPU를 삭제합니다.

    TPU VM

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

    TPU 노드

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

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

    TPU VM

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

    TPU 노드

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a

다음 단계

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