1. TPU가 뭐야?

Tensor Processing Unit의 줄임말이다. GPU는 그래픽 처리나 여러 병렬처리 연산에 사용되는 범용적인 프로세서이고 딥러닝에서도 GPU를 이용한 분산 처리가 일반적이다. TPU는 딥러닝에서 가중치를 계산하는, 즉 tensor를 여러 코어가 분산해서 동시에 계산하는데 특화된 하드웨어 가속기이며 구글에서 개발하고 지원 중이다.
Kaggle에서 무료로 이용할 수 있다. 주당 20시간, 그리고 한 세션에 최대 9시간 이용 가능하다.

2. TPU 사용법

Kaggle notebook을 이용한다면 Accelarator를 TPU v3-8로 설정하는 것이 시작이다. 화면의 우측 상단에서 다음과 같이 변경한 후 세션을 요청하면 된다.
참고로 GPU T4 x2로 변경하면 세션을 바로 생성되는 것과 다르게, 사용자가 많으면 대기를 해야한다.

이후에 아래 코드를 이용해서 tpu를 이용하면 된다.

# detect and init the TPU
tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()

# instantiate a distribution strategy
tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)

# instantiating the model in the strategy scope creates the model on the TPU
with tpu_strategy.scope():
    model = tf.keras.Sequential( … ) # define your model normally
    model.compile( … )

# train model normally
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=…)

위 코드는 kaggle에서 제공하는 TPU 소개글에 있는 코드이다. 주석으로 이해해볼 수 있지만 정리를 해보면 아래와 같다.
1. TPU 서버에서 TPU Cluster를 요청하고 반환 받는다.
2. 할당 받은 TPU 유닛을 이용해서 strategy를 정의한다.
3. strategy.scope() 메소드 범위 안에서 model을 정의하고 컴파일한다.
4. 학습시킨다. 이때 strategy.scope() 범위를 벗어나서 진행하면 된다.

3. batch size & learning rate 설정

3-1. batch size

공식 문서에 따르면 TPU v3-8의 경우 코어가 8개이고, 경험상으로는 하나의 코어 당 128개의 element를 돌리는 걸 추천한다고 한다. 이 정도 규모에서 일반적으로 코어가 쉬지 않고 돌아간다고 하고, 코어 당 8개 정도로 계산해서 배치 사이즈를 설정해도 속도가 향상되는 걸 체감할 수 있다고 한다.

BATCH_SIZE = 16 * strategy.num_replicas_in_sync

여기서 strategy.num_replicas_in_sync는 현재 이용 가능한 tpu 유닛의 수를 나타낸다고 한다.

개인적인 사용 경험으로는, 당연하게도 처리하는 데이터 하나의 크기에 따라 처리 가능한 사이즈로 조절을 해주어야 한다.

3-2. learning rate

배치 사이즈가 커짐에 따라 learning rate도 커져야 한다. 공식 문서에서 이 부분 역시 경험적으로 그렇다고 한다.
추가적으로 learning rate를 조절하는 것은 Bag of Tricks(2018) 논문을 참고하면 도움이 된다. 간단히 말하면 배치 사이즈를 키울 때 특정 수준을 넘어가면 에러가 급격하게 늘어나는데, learning rate warm-up과 learning rate decay라는 방법을 섞어서 사용하면 에러가 급격하게 늘어나는 지점을 더 키울 수 있고, 이 말은 에러를 늘리지 않는 선에서 큰 배치 사이즈를 확보할 수 있게 된다.

4. 코드를 보면서 생기는 궁금증

4-1. tf.distribue가 뭐야?

• tensorflow로 분산처리를 할 때 쓰는 패키지이다.

4-2. resolver는 뭐고 tf.distribute.cluster_resolver는 뭐야?

• 프로그래밍에서 resolver는 (gpt에게 물어보니) 어떤 요청을 받아들이고 그에 맞는 응답을 알려주는 객체 정도로 이해하면 된다. 일반적으로 데이터베이스나 클라이언트-서버 구조를 만들어서 다룰 때 resolver라는 말을 보게 되는데, 원하는 쿼리나 요청을 보내면 적절한 응답을 받는, 일종의 네트워크를 이용한 통신을 대신해주는 역할을 한다고 생각하면 된다.
• 같은 의미로 cluster_resolver의미를 풀어서 이해해보면, TPU cluster를 요청하고 받아서 프로그래머에게 쥐어주는 역할을 한다고 볼 수 있다. TPU는 GPU처럼 독립된 상태의 하드웨어로 개인이 사서 컴퓨터에 붙이는 장비가 아니고 구글에서 관리하는 서버에서 자원을 논리적으로 요청/할당 받아서 사용하게 된다. 때문에 이용하려면 쓰려는 TPU 자원을 할당받아서 써야하기 때문에 resolver가 존재하는 것이다.

4-3. tf.distribute.experimental은 뭐야?

• tensorflow에서 지원은 하지만 안정성이 보장되지 않는, 말 그대로 실험적인 기능인 것이다. 카카오톡에 있는 실험실이나 크롬에서 지원하는 실험 기능 같은 느낌이다.
• 분산 처리할 때 Strategy라는 모듈을 이용하게 되는데, TPUStrategy는 experimental의 하위 모듈이다.