파이썬 코드로 GPU memory print 하기

import pynvml
import torch
def get_memory_free_MiB(gpu_index):
    pynvml.nvmlInit()
    handle = pynvml.nvmlDeviceGetHandleByIndex(int(gpu_index))
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    total= mem_info.total // 1024 ** 2
    free = mem_info.free // 1024 ** 2
    return free, total

print('GPU Memory : ',{f'GPU-{i}':f'{get_memory_free_MiB(i)[0]:,d}/{get_memory_free_MiB(i)[1]:,d} MiB' for i in range(torch.cuda.device_count())})

 

 

print('GPU Memory : ',{f'GPU-{i}':f'{get_memory_free_MiB(i)[0]:,d}/{get_memory_free_MiB(i)[1]:,d} MiB' for i in range(torch.cuda.device_count())})

 

or

print('GPU Memory : ',{f'GPU{i}':f'{get_memory_free_MiB(i)[0]:,d}/{get_memory_free_MiB(i)[1]:,d} MiB' for i in os.environ['CUDA_VISIBLE_DEVICES'].split(',')})