딥러닝 모델을 학습시키다 보면 항상 vram의 압박에 시달리게 된다. 특히 최근 막대한 크기의 모델들이 등장해 이런 압박은 더 심해지기도 한다.
한편, 일반 사용자용 그래픽 카드 중 최상위인 Nvidia 2080ti조차도 vram이 겨우 11GB밖에 되지 않아 거대한 모델을 Fine-tuning 하는 것조차 굉장히 작은 배치사이즈로 학습시켜야 한다.
Google Colab에서 제공하는 TPU는 tpu v3-8 모델로 총 128GB의 메모리를 가지고 있어, 상대적으로 큰 모델과 배치사이즈를 이용해 학습할 수 있다. (tpu v3 하나는 16GB의 HBM 메모리를 가지고 있고, tpu v3-8은 8개의 코어로 총 128GB의 메모리를 가진다.)
PyTorch에서는 Pytorch/XLA 프로젝트를 통해 PyTorch에서도 TPU를 통한 학습을 할 수 있도록 컴파일러를 제공하고 있고, colab에 해당 패키지를 설치하면 TPU를 곧바로 사용할 수 있다.
NOTE: 이번 글은 아래 공식 튜토리얼의 내용을 따라갑니다.
공식 Tutorial: PyTorch on Cloud TPUs: Single Core Training AlexNet on Fashion MNIST
(단 내용의 100%를 담는 대신, 기존 PyTorch와 동일한 부분은 제외함)
