아래 추가적인 코드 설정을 해주지 않으시면 실제로는 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)
tf.distribute.Strategy는 훈련을 여러 GPU 또는 여러 장비, 여러 TPU로 나누어 처리하기 위한 텐서플로 API입니다. 이 API를 사용하면 기존의 모델이나 훈련 코드를 분산처리를 할 수 있습니다. TPU 사용을 위해서도 Strategy를 셋팅해줍니다.
strategy = tf.distribute.TPUStrategy(resolver)
딥 러닝 모델을 컴파일을 할 때도 추가적인 코드가 필요
모델의 컴파일은 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를 사용하며 학습