Colab에서 TPU로 BERT 처음부터 학습시키기 - Tensorflow/Google ver.

2018년말부터 현재까지 NLP 연구에서 BERT는 여전히 압도적인 위치를 차지하고 있다.

한편, BERT모델을 사용하는 이유 중 가장 큰 것 하나가 바로 한국어로 Pretrained된 모델이 있다는 점이다. Google에서 논문을 처음 공개했을 때 Multilingual pretrained model을 공개해 Fine-tuning만으로도 우리가 필요한 데이터셋에 맞춰 분류기를 만드는 등의 여러 응용이 가능하고, 동시에 높은 성능을 보여주었기 때문에 BERT 자체를 학습시키는 것에 대해서는 크게 신경쓰지 않은 것이 사실이다.

한편 작년 ETRI의 한국어 BERT 언어모델, 그리고 SKTBrain의 KoBERT 등 한국어 데이터셋으로 학습시킨 모델들이 등장했고, 이런 모델들을 Fine-tuning할 경우 기존 구글의 다국어 모델을 사용한 것보다 성능이 조금이라도 더 잘 나오기도 한다. (특히 정제되지 않은 글에 대해 좀 더 나은 성능을 보여줬다. OOV문제가 덜한 편이었다.)

다만 이런 모델들 역시 굉장히 ‘보편적’ 글로 학습된 것이라 도메인 특화된 분야에 있어서는 성능이 잘 나오지 않을 수도 있다. 따라서 특수한 경우의 특수한 도메인에 최적화된 Pretrained model을 만든다면 우리의 NLP 모델도 좀 더 성능이 좋아질 수 있다!

이번 글에서는 BERT 모델을 TPU와 Tensorflow를 이용해 처음부터 학습시켜보는 과정을 다뤄본다.

이번 글은 Colab Notebook: Pre-training BERT from scratch with cloud TPU를 기반으로 작성되었습니다.

Colab에서 PyTorch 모델 TPU로 학습하기

딥러닝 모델을 학습시키다 보면 항상 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와 동일한 부분은 제외함)

Your browser is out-of-date!

Update your browser to view this website correctly. Update my browser now

×