Pytorch 특정 GPU 사용하기

나는 많은 딥러닝 프레임워크 중 Pytorch와 MxNet을 자주 사용하는 편이다.
그런데, 연구실 사람들과 GPU 서버를 함께 쓰다 보니 어떤 GPU가 현재 available한지 알아야 할 필요가 있었다.

원래는 시간대를 적당히 맞춰서 사용하곤 했는데, 멀티 GPU 세팅에 대해서는 잘 모르기도 하고 colab으로는 도저히 불편해서 쓸수가 없었으므로..ㅠㅠ 관련 정보를 찾아보고 그 결과를 간단히 공유하고자 한다.

GPU 서버 상태 확인하기

사용하고자 하는 서버의 GPU availability 를 먼저 알아보자. 간단하게 알아볼 수 있다.

1 2	nvidia-smi # bash OR terminal에서 명령어를 입력할 경우 !nvidia-smi # Jupyter Notebook에서 명령어를 입력할 경우

위와 같이 입력한 경우, 아래와 같은 화면이 출력된다.
(만약 출력이 안되는 경우, Nvidia driver에 문제가 있는 것이기 때문에 요 부분을 다시 검토해 보자)

내가 사용하는 서버는 TITAN Xp 가 4개 가동중인 GPU 서버이며, 현재 1,2,3번 GPU에서는 용량이 꽉 차있는 것으로 보아 누군가가 작업을 진행하고 있는 것으로 예상된다.

Default GPU Setup 확인하기

어떤 GPU를 사용할지는 위의 과정을 거쳤더니 알게 되었다. 나는 GPU:0 을 사용해야 하는구만.

그러면 이제 내 작업환경으로 돌아가, 현재는 어떤 GPU를 사용하도록 셋팅 되어있는지 확인하는 과정이 필요하다.
마치 os.getcwd() 으로 현재 디렉토리를 아는 것처럼..

import torch

# 현재 Setup 되어있는 device 확인
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print ('Available devices ', torch.cuda.device_count())
print ('Current cuda device ', torch.cuda.current_device())
print(torch.cuda.get_device_name(device))

위 코드를 실행시키면 차례대로 사용가능한 GPU 디바이스의 갯수, 현재 셋업 되어있는 GPU 넘버, 그리고 GPU 디바이스의 이름을 출력한다.

결과를 보니, 나는 TITAN Xp 디바이스 중 2번 GPU를 사용중임을 확인할 수 있었다. (인덱스는 0번부터 시작)

GPU Allocation 변경하기

나는 현재 2번 GPU를 사용중인데, 이를 0번 GPU로 바꾸어야 한다.

# GPU 할당 변경하기
GPU_NUM = 0 # 원하는 GPU 번호 입력
device = torch.device(f'cuda:{GPU_NUM}' if torch.cuda.is_available() else 'cpu')
torch.cuda.set_device(device) # change allocation of current GPU
print ('Current cuda device ', torch.cuda.current_device()) # check

# Additional Infos
if device.type == 'cuda':
    print(torch.cuda.get_device_name(GPU_NUM))
    print('Memory Usage:')
    print('Allocated:', round(torch.cuda.memory_allocated(GPU_NUM)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_cached(GPU_NUM)/1024**3,1), 'GB')

위 코드를 실행시키면 0번 GPU를 사용하도록 변경하고, 해당 GPU의 cuda memory 상태를 간략하게 출력한다.

제일 위의 출력 결과물을 확인해 보니, 성공적으로 Current cuda device가 2에서 0으로 바뀌었음을 알 수 있다.

모쪼록 GPU 서버가지고 싸우지 말고(?) 딥러닝을 즐기도록 하자 :)