Colab TPU 설정

정강민·2022년 1월 31일
0

AI

목록 보기
3/3

  1. 코랩(Colab)에서 TPU를 선택
    Colab에서 GPU를 사용하기 위해서 런타임 유형 변경에서 GPU를 선택하였듯이 TPU를 사용하기 위해서는 TPU를 선택
  • Colab에서 런타임 > 런타임 유형 변경 > 하드웨어 가속기에서 'TPU' 선택

    아래 추가적인 코드 설정을 해주지 않으시면 실제로는 TPU를 사용하지 않으므로 아래의 코드 설정을 필수!

  1. TPU 초기화
import tensorflow as tf
import os

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])

tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
  1. TPU Strategy 셋팅

tf.distribute.Strategy는 훈련을 여러 GPU 또는 여러 장비, 여러 TPU로 나누어 처리하기 위한 텐서플로 API입니다. 이 API를 사용하면 기존의 모델이나 훈련 코드를 분산처리를 할 수 있습니다. TPU 사용을 위해서도 Strategy를 셋팅해줍니다.

strategy = tf.distribute.TPUStrategy(resolver)
  1. 딥 러닝 모델의 컴파일

딥 러닝 모델을 컴파일을 할 때도 추가적인 코드가 필요
모델의 컴파일은 strategy.scope 내에서 이루어져야 합니다. 이는 모델의 층을 쌓는 create_model()와 같은 함수를 만들어 strategy.scope 내에서 해당 함수를 호출하여 모델을 컴파일하는 방식사용

다음과 같이 모델의 층을 쌓는 create_model()라는 함수

def create_model():
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),
       tf.keras.layers.Conv2D(256, 3, activation='relu'),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(256, activation='relu'),
       tf.keras.layers.Dense(128, activation='relu'),
       tf.keras.layers.Dense(10)])

with strategy.scope(): 다음에 들여쓰기를 하고, create_model() 함수를 호출하고 모델을 컴파일합니다. 다시 말해 strategy.scope 내에서 모델을 컴파일

이 모델을 fit() 하게되면 해당 모델은 TPU를 사용하며 학습

profile
DA/DA/AE

0개의 댓글