AI모델 학습시 Batch Size와 학습 속도의 상관관계

 
  AI 분야에서 관련 연구를 수행하시는 분들이라면 그 누구도 Nvidia의 GPU를 사용하지 않고 AI모델을 학습을 할 수 있는 방법을 떠올리기 어려울 정도로 Nvidia GPU는 2024년 현재 인공지능 분야에서 대체불가의 절대 권좌를 누리고 있습니다. 그러한 영향으로 Nvidia의 주식은 2024년초 대비 무려 2배 가까이 상승하였고, AI 적용 분야의 확산으로 인한 GPU의 수요는 끊임없이 상승할 것으로 보여 Nvidia의 전망을 매우 밝아보입니다.
 AI 연구를 수행하고 있는 저 또한 여러 종류의 Nvidia GPU를 접할 수 있는 기회가 있었습니다. GTX 1080 시리즈부터 H100까지 지금도 다양한 그래픽카드를 사용하여 여러 종류의 AI모델들의 학습을 진행하는 연구를 수행하고 있습니다. 워낙 많은 데이터들을 다루다보니 결과 분석을 위해 학습 수행 속도 향상 방법 또한 상당히 많은 고민을 하고 있습니다. 지금까지 그래픽카드들을 사용함에 있어, 학습 속도를 향상시키기 위해 저는 다음과 같은 사항들을 고려해왔습니다.
 

1. 실제 사용 환경에서 최적화된(너무 크지도, 너무 작지 않은) AI모델의 Weight Size
2. GPU 메모리 크기 및 연산 속도
3. AI 모델 학습 수행시 Batch Size의 크기
 
GPU 메모리 크기의 경우 하드웨어적인 성능에 해당되므로 좋은 성능의 GPU를 사용하였을 때 모델의 학습속도가 향상됨은 분명합니다. AI모델 Size의 경우 자신이 학습하고자 하는 데이터의 특성에 맞추어 적정한 크기의 모델을 선정하는 과정으로 너무 큰 사이즈의 AI모델 선정시 학습이 완료되는데 소요시간이 상당히 길어지며, 너무 작은 사이즈의 AI모델 선정시 학습 소요시간이 줄어드나 성능이 나오지 않는 경우가 발생하기 때문에 자신의 데이터셋 환경에 적합한 최적의 크기를 자세히 따져볼 필요가 있습니다.
 
 보편적으로 Batch Size의 경우 클 수록 AI모델의 성능 및 학습속도가 향상되는 것으로 알려져 있고 실제 학습 결과를 통해 성능을 확인할 수 있습니다.
 AI모델의 성능 관점에서 보았을 때, Batch Size가 클 수록 개별 데이터로 학습할 때보다 각 데이터별 편차로 인한 Noize를 줄일 수 있고 Testset 결과에서 정답에 좀 더 가깝게 수렴할 수 있습니다. 다만, 너무 큰 Batch Size를 사용할 경우 Overfitting이 발생하여 AI모델이 Local Optima에 도달할 수 있다는 단점이 있습니다.
 학습속도의 관점에서 보았을 때, 개별 Data를 한 번씩 학습하는 것보다 Batch를 통해 여러개의 Data를 Batch Normalization을 거친 후 AI모델의 학습 입력으로 사용하기 때문에 AI모델의 학습 횟수가 (DatasetSize/BatchSize)만큼 줄어들어 AI모델이 학습하는데 걸리는 시간을 단축할 수 있습니다.
 

 
 위 그림을 보았을 때 총 128개의 데이터로 구성된 데이터셋으로 학습 수행시 Batch사이즈가 2인경우 64회, Batch사이즈가 4인 경우 32회의 학습으로 데이터셋의 전체 데이터를  학습할 수 있습니다. 즉, Batch를 사용하여 학습 수행시 학습 시간을 단축시키면서 AI모델의 성능 향상 또한 가능하다는 것이지요.

https://arxiv.org/abs/2105.07576

 메타리서치(구 Facebook AI Research)에서 2021년 투고한 논문에서 Batch Size가 학습에 미치는 영향을 다루는 연구 내용을 다루었는데, 실제 Batch Size가 커질수록 Error Rate가 줄어들고, Train Dataset과 Test Dataset의 결과 차이가 줄어드는 것을 확인하였다고 합니다. 논문에 따르면, Batch Size가 256-512 사이일 때 Train-Test 수행시 가장 효율적이었다고 합니다.
 저 또한 지금까지 AI모델을 개발한 경험들을 통해 저는 다음과 같은 생각을 하였습니다.

GPU 메모리 용량이 크다 = Batch Size를 크게 설정할 수 있다 = 학습 속도가 빨라진다?

 그런데 이러한 저의 생각을 무참히 부수어버린 케이스가 발생하였습니다. 연구소내 서버에 Nvidia H100(80GB)이 도입되어 이를 가지고 학습할 수 있는 기회가 생겨 기존에 쓰던 RTX 4090(24GB) 환경에서 학습할때 사용하였던 AI모델을 사용해서 기존에 학습하던 환경에서보다 Batch Size를 두배 늘려서 학습을 수행해보았습니다.
 그렇게 한껏 기대에 차있던 상황에서 학습 과정을 지켜보던 저는 당혹스러운 상황을 마주치게 되었습니다. 기존 환경에서 학습하였을 때와 소요시간이 별 차이가 없던 것입니다. 기존에 사용하던 GPU보다 더 좋은 성능의 환경에서 놀라운 결과를 기대했던 저에게 놀라움과 실망감이 교차하는 순간이었습니다. 

  연구를 수행하는 직종에 종사하는 사람의 입장으로서 '왜 이러한 현상이 발생하는 것인가?'라는 원초적인 질문을 던지고 원인을 찾아보기로 하였습니다. 관련 문제들에 대해 검색을 하다보니 GTX 1080 Ti와 Titan V, GV100 그래픽카드를 가지고 Batch Size와 학습 속도를 다루었던 내용을 찾을 수 있었습니다.
 

pugetsystems.com

  위 그래프를 보았을 때,  전반적으로 Batch Size가 커질수록 초당 이미지 처리 갯수가 많아지는 경향이 있는 것을 알 수 있습니다. 다만 Batch Size가 두 배로 커진다고 해서 학습 시간이 절반으로 줄어들지 않으며, Batch Size가 너무 커지면 오히려 학습 속도가 떨어지는 경우가 발생합니다. Batch Size의 관점에서 보았을 때는 GPU 메모리 크기가 학습 속도에 영향을 줄 것으로 보이지만 또다른 요소들로 인해 AI모델의 학습 속도에 영향을 주기 때문에 Batch Size를 마냥 키우는것만이 만능은 아닙니다. 왜냐하면 자신이 사용하는 PC의 환경 또한 학습 속도에 영향을 주기 때문입니다. 학습 속도에 영향을 주는 요인으로 다음과 같은 사항들이 있습니다.

·GPU 메모리 데이터

 AI모델 학습 수행시 GPU 메모리에는 Batch 데이터 뿐만 아니라 학습이 수행되는 AI모델 또한 GPU 메모리에 올라갑니다. Weight Size가 큰 AI모델이 GPU에 올라갈 경우 AI모델이 GPU 메모리 전체에서 차지하는 비율이 높아지므로, 나머지 공간을 사용하게 되는 Batch 데이터의 비중이 축소될 수 밖에 업습니다. 즉, GPU 메모리의 활용 용량 한계로 인해 학습 수행시 Batch Size를 필연적으로 줄여야 할 수 밖에 없는 것입니다.

·오버헤드

  컴퓨터의 RAM, CPU, 하드디스크 등의 장치 성능이 AI모델의 학습 속도에 영향을 줄 수 있습니다. 그래픽카드에서 학습이 진행되기 위해 CPU가 하드디스크로부터 Batch 데이터를 GPU 메모리로 갖고 오도록 명령하고, Batch 데이터를 하드디스크로부터 GPU에 전달되기 까지의 과정에서 소요되는 시간을 오버헤드(Overhead)라고 표현합니다. 즉, 성능 좋은 GPU를 사용한다 하더라도 하드웨어의 스팩으로 인한 소요시간 또한 상당히 발생할 수 있는 것입니다.

·멀티GPU

 CPU가 성능 대비 전력소모량이 기하급수적으로 증가하는 현상으로 인해 멀티코어CPU로 발전하여 병렬 다중 처리 방식으로 발전했듯이 GPU 또한 여러 대의 GPU를 사용하여 AI모델 학습을 하는 사례들이 많습니다. 한때 암호화폐 채굴을 위해 PC 한 대에 여러 대의 GPU를 사용하여 채굴 하는 등의 사례로 멀티GPU를 경험하신 분들도 계실 정도로 주변에서 멀티GPU를 사용하는 사례도 많아 접할 수 있습니다. 다만, Pytorch로 멀티GPU 연산 수행시 한동안 Windows 환경에서 멀티GPU를 사용하는데 제약이 많아 Linux 운영체제를 사용하는 것이 필연적이었지만, Windows11부터 WSL2 환경에서 Linux를 사용할 수 있게 됨으로서 Windows 운영체제에서도 멀티GPU 개발을 기존보다 더욱 쉽게 할 수 있게 되었습니다.

 멀티GPU로 AI학습을 수행할 경우 각 GPU마다 AI모델 및 Batch가 올라가게 됩니다. 예를 들어, 2대의 GPU를 사용하였을 때, 2대의 GPU에 똑같은 Weight 구조의 AI모델이 각각 올라가고, Batch 또한 전체 Batch를 절반씩 나누어 각각의 GPU에 올라가게 됩니다. 즉, 총 64개의 Batch가 있을 경우 각 GPU별로 32개의 Batch가 GPU에서 학습이 진행되는 것입니다.

 멀티GPU 환경에서 학습 수행시 GPU가 많아질수록 Batch가 각 GPU에 분산되어 학습을 수행하므로 자연스럽게 전체 학습 속도의 단축 효과를 확인하실 수 있습니다. 다만, 모든 GPU 메모리에 올라가는 데이터의 구조가 동일하기 때문에(동일한 크기의 AI가 각각의 GPU에 올라감) Batch Size의 크기는 단일GPU에서 사용할 때와 차이가 없는 제약이 있습니다. 멀티GPU 학습시 Batch가 어떻게 학습되는지 혹시 궁금하신 분들은 제가 이전에 작성하였던 Synchronized Batch Normalization 글을 참조해주시기 바랍니다.

 

2024.02.12 - [AI용어정리] - [용어정리]Synchronized Batch Normalization

[용어정리]Synchronized Batch Normalization

 결론 

 짧게나마 Batch Size가 AI모델의 학습 속도에 미치는 영향과 총학습시간을 단축시킬 수 있는 방법들에 대해 이야기를 다루어보았습니다. 지금까지 설명한 내용을 기반으로 AI모델의 학습속도를 향상시킬 수 있는 방법으로 저는 다음과 같은 방법들을 제시해봅니다.

 - 학습 데이터셋에 최적화된 AI모델 Weight Size 설정

 - 최대한 큰 Batch Size를 설정할 수 있는 GPU 확보

 - Train 진행시 오버헤드를 고려해 GPU 메모리 점유 비중(AI모델+Batch Size) 을 50-70% 사이에서 설정 

 - 같은 모델의 GPU 여러대로 멀티GPU 환경 구축
 

참고자료:
WU, Yuxin; JOHNSON, Justin. Rethinking" batch" in batchnorm. arXiv preprint arXiv:2105.07576, 2021.

Rethinking "Batch" in BatchNorm

 

https://www.pugetsystems.com/labs/hpc/gpu-memory-size-and-deep-learning-performance-batch-size-12gb-vs-32gb-1080ti-vs-titan-v-vs-gv100-1146/

GPU Memory Size and Deep Learning Performance (batch size) 12GB vs 32GB -- 1080Ti vs Titan V vs GV100