Cloud TPU 사용하기

@아래 내용은 TPU quota 가 신청되어 있어야 사용 가능한 부분입니다. 참고 하시기 바랍니다.

TPU 가 일반에 공개된 것과 비슷한 시기에 Google Compute Engine (GCE) 에도 TPU 메뉴가 생성되었다. 와우, 이제 우리의 인스턴스 내에서 TensorFlow 를 통해 기계학습을 진행할 때 TPU 를 지정해서 사용할 수 있게 된 것이다. 이 부분은 CloudML 과는 전혀 상관이 없으므로 오해하지 않도록 한다. Google Cloud Platform (GCP) 의 모든 기능이 그러하듯이 Cloud TPU 도 API 사용 설정을 통해 기능을 활성화시켜야 사용할 수 있다.

 기능을 활성화하면 잠시 후 아래와 같은 메뉴를 볼 수 있다. 일단 TPU 노드를 생성해보자. (빠른 시작 사용을 통해 기능 사용법의 상세를 보면 더욱 좋다. 콘솔에서 활용하는 방법에 대한 설정이 자세하게 되어 있다)

설정은 다소 심플하다. 아래처럼 TPU 를 사용할 리전과 TensorFlow 버전, 네트워크 등을 설정해야 한다. 선택의 폭이 크지 않기 때문에 고민할 것도 없다.

그리고 여기서 가격을 눈 크게 뜨고 확인해보시라. 시간당 $6.50, 730 시간 기준으로 $4,745 다. 언뜻보면 너무 비싸서 거부감이 들 수도 있는데 일반 GPU 로 24시간 학습시켜야 할 네트워크를 훨씬 더 빠른 시간안에 학습시킬 수 있기 때문에 비용적으로는 사실 더 이득을 취할 수 있겠다. 또한 K80, P100 GPU 의 한달 임대 비용도 몇천 달러 정도는 나간다. 아무튼 적지 않은 가격이기 때문에 사용할 때만 노드를 생성하고 사용이 끝나면 노드를 삭제하는 습관이 꼭 필요할 듯 싶다.

네트워크 설정 부분은 실제 기계학습 코드를 구현할 때 TensorFlow 세션을 TPU 쪽으로 열때 사용된다. IP 주소 범위를 10.0.101.0 으로 지정한다면 TPU_IP 는 10.0.101.2 가 된다는 점을 참고하도록 하자. 구현한다면 대략 아래와 같은 코드가 된다. (공식 예제)

import os
import tensorflow as tf
from tensorflow.contrib import tpu
def computation(alpha, x, y):
  return alpha * x + y
alpha = tf.Variable(3.0, name='alpha')
x = tf.Variable(tf.ones([3, 3], tf.float32), name='x')
y = tf.Variable(tf.ones([3, 3], tf.float32), name='y')
tpu_computation = tpu.rewrite(computation, [alpha, x, y])
with tf.Session('grpc://{TPU_IP}:8470'.format(**os.environ)) as sess:
  sess.run(tpu.initialize_system())
  sess.run(tf.global_variables_initializer())
  output = sess.run(tpu_computation)
  print(output)
  sess.run(tpu.shutdown_system())

  print 'Done!'

이제 GCE 에서 TPU 를 사용해서 기계학습을 돌려보자. 그리고 GPU 와 속도차이를 느껴보면 되겠다. 그리고 CloudML 에서 TPU 사용하려면 아직은 신청을 통해 선정되는 듯하니 참고하도록 하자. 이번 업데이트를 통해 기계학습 유저가 GCE 로 많이 돌아가려나? 아무튼 퍼블릭 클라우드의 점유율은 AWS, MS 에 다소 밀리지만 기계학습쪽으로는 명실상부 최고가 아닌가 싶다. 이번 TPU 공개로 쐐기를 박을 듯.