nvidia_nvml_gpu_memory_usage_and_driver_info.py

1import nvidia_ml_py3.nvml as nvml
2
3nvml.nvmlInit()
4try:
5    deviceCount = nvml.nvmlDeviceGetCount()
6    print(f"Driver Version: {nvml.nvmlSystemGetDriverVersion().decode('utf-8')}")
7    
8    for i in range(deviceCount):
9        handle = nvml.nvmlDeviceGetHandleByIndex(i)
10        info = nvml.nvmlDeviceGetMemoryInfo(handle)
11        print(f"Device {i}: {nvml.nvmlDeviceGetName(handle).decode('utf-8')}")
12        print(f"  Memory Total: {info.total / 1024**2:.0f} MB")
13        print(f"  Memory Used: {info.used / 1024**2:.0f} MB")
14        print(f"  Memory Free: {info.free / 1024**2:.0f} MB")
15finally:
16    nvml.nvmlShutdown()