[PyTorch] 리눅스환경에서 특정 GPU만 이용해 Multi GPU로 학습하기

 

 

연구실에서 다같이 연구를 하다보니 같은 서버에 GPU를 나눠서 사용해야 한다. 그럴 때 어떻게 나눠쓸 수 있을까? CUDA_VISIBLE_DEVICES를 이용해서 내가 현재 볼 수 있는 GPU를 한정시키는 방법이 있다. 그래서 특정 GPU들만 이용해서 Multi GPU로 모델을 학습할 수 있다. 예를 들면, 만약 2,3번만 보이게 설정했다면 걔네를 GPU 0번과 1번으로 인식하여 두 GPU만을 이용해서 학습하게 된다.

 

1) 터미널 창에서 GPU를 지정하여 python파일을 실행시킨다.

 

model = torch.nn.DataParallel(model)
model.cuda()

 

CUDA_VISIBLE_DEVICES = 2,3 python train.py #터미널창

 

2) 코드내부에서 GPU를 지정해서 실행시킨다.

 

import os
os.environ["CUDA_VISIBLE_DEVICES"]='2, 3'

model = torch.nn.DataParallel(model)
model.cuda()

 

python train.py #터미널창