Multi-gpu 환경에서 자동으로 적절한 GPU 하나 골라서 학습하기


하나보다 많은 GPU를 가지고 있더라도 언제나 딥러닝 모델을 multi-gpu로 데이터 병렬화 혹은 모델 병렬화를 구현하여 학습을 진행하는 것은 쉽지 않은 일이다. PyTorch처럼 처음부터 다양한 구현 모델을 염두해두었거나, 오픈소스 커뮤니티가 최근 가장 활성화되어 있는 라이브러리에서는 병렬화 구현이 유리하지만, TensorFlow 같은 경우 긴 세월에 걸친 버젼 업그레이드와 수시 때때로 deprecated 되는 API로 인하여 코드 유지보수, 버져닝, 호환성 개선, API의 버젼에 따른 병렬화 방법론의 차이를 보여 API 문서를 제법 잘 따라하더라도 모델에 따라서는, 혹은 전체 구현이 너무나 single-gpu를 기준으로 만들어져 효율적인 multi-gpu 병렬 학습의 구현이 어려운 경우도 있다. 물론 시간이 많거나 다른 제약이 없다면 정해면 프레임워크 내에서도 충분히 병렬화를 구현할 수 있지만, 그럴 수 없다면 모델이나 데이터 병렬화가 아닌 ‘실험’의 병렬화도 자원을 최대로 활용하는 방법이다.

여건이 어려워 실험을 병렬화하는 것 같지만, 실험의 병렬화에는 생각지도 못한 엄청난 장점이 있다. business hour 측면에서 single-gpu로 3일이면 학습이 끝나는 모델 A가 있다고 한다면, 이 모델을 금요일에 구현하여 학습을 시작하였을 때 4-gpu를 사용한다고 하더라도 학습은 빨라도 하루정도는 걸리기 때문에 월요일에 결과를 확인해야 하지만, 실험 병렬화를 구현할 경우 4-gpu마다 각각의 실험 설정의 학습을 진행하여 4개의 모델의 학습 결과를 월요일 한 번에 확인할 수 있다. 이는 데이터셋이 비교적 작은 특정 도메인에 한정적인 실험에 굉장히 유리하다.

실험 병렬화 시나리오

  • 하나로로 학습, 하나로는 검증 및 디버깅, 여타 소프트웨어에 물려서 개발
    • GPU 0: Training
    • GPU 1: Evaluation / Debugging
  • 여러 모델 학습
    • GPU 0: Model A
    • GPU 1: Model B
    • GPU 2: Model A with hparam B
    • GPU 3: Model B with hparam A

GPUtil을 통해 자동으로 NVIDIA GPU 하나 고르기

GPUtil 설치

$ pip install gputil

GPU 선택 자동화

./train.py

import os
import GPUtil


def choose_one_gpu():
    """set CUDA_VISIBLE_DEVICES environment variable to utilize one best gpu """
    os.environ["CUDA_VISIBLE_DEVICES"] = str(GPUtil.getFirstAvailable()[0])


def train():
    choose_one_gpu()
    some_train_ops(...)


if __name__ == "__main__":
    train()

실험 병렬화

# tmux, byobu 등 터미널 멀티플렉서에서 실행하면 관리에 용이하다.
# CAUTION: 실험 코드를 동시에 실행할 경우 동일 GPU를 점유하려고 시도해 충돌될 수 있다.
# 앞서 실행한 실험이 gpu allocation을 진행한 경우에만 안정적으로 GPU가 자동할당된다.
python train.py experiment_01.yaml
python train.py experiment_02.yaml
python train.py experiment_03.yaml
python train.py experiment_04.yaml