Python 내에서 GPU 메모리 확인하는 법

소위 말해 python 안에서 nvidia-smi 를 수행하는 방법이다.

pip install nvidia-ml-py3

위 패키지를 설치하고 아래와 같이 패키지를 불러와 사용한다.

import nvidia_smi

nvidia_smi.nvmlInit()
deviceCount = nvidia_smi.nvmlDeviceGetCount()
freem = []
for i in range(deviceCount):
    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
    info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
    freem.append(100*info.free/info.total)
nvidia_smi.nvmlShutdown()

 

이제 freem 리스트에 각 GPU의 메모리 사용량이 % 단위로 담기게 된다.

 

아래의 코드를 추가해서 자동으로 가장 여유로운 GPU를 선택해 사용하도록 응용할 수 있다.

import torch

device_num = np.argmax(freem).item()
torch.cuda.set_device(device_num)
device = torch.device(f"cuda:{device_num}")

이후 모델이나 텐서에 .to(device)를 통해 해당 GPU에 할당하면 된다.