GPUのメモリ確認

GPUのメモリを確認する方法を説明します。

インストール

pip install nvidia-ml-py3

確認

import nvidia_smi

nvidia_smi.nvmlInit()

deviceCount = nvidia_smi.nvmlDeviceGetCount()
for i in range(deviceCount):
    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
    info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
    print("Device {}: {}, Memory : ({:.2f}% free): {}(total), {} (free), {} (used)".format(i, nvidia_smi.nvmlDeviceGetName(handle), 100*info.free/info.total, info.total, info.free, info.used))

nvidia_smi.nvmlShutdown()

Device 0: b'Tesla T4', Memory : (98.32% free): 16106127360(total), 15835267072 (free), 270860288 (used)

参考