pytorch에서 gpu 슬롯 별 메모리 여유량 알아보는 기능

1. 개요

gpu사용량을 파이썬 내에서 실시간으로 알아내면 모델 학습시에 편리하기 때문에 아래의 기능을 만듬.

먼저 print기능을 통해 gpu 슬롯 별 여유량을 출력하게 만듬.

그 후에 파이썬의 리스트에 각 슬롯번호 순서대로 메모리 사용량을 저장하게 만듬.

import torch
def get_gpu_memory_usage():
    memories=[]
    for i in range(torch.cuda.device_count()):
        free, total = torch.cuda.mem_get_info(f'cuda:{i}')
        mem_used_mb = (total - free) / 1024 ** 2
        print(f'cuda:{i} free :',free/1024**3,'Gb')
        memories.append(mem_used_mb)
    return memories

 

2. 결과물

위의 기능을 실행한 결과

위와 같이 각 gpu 슬롯 별로 여유공간이 출력됨을 볼 수 있음.