Print out CPU/GPU memory usage (basic)

get_memory_usage.py

import math
import psutil
from pynvml import nvmlDeviceGetCount, nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo


def format_bytes(nbytes: int) -> str:
    if nbytes == 0:
        return "0 B"

    unit = ("B", "kB", "MB", "GB", "TB")
    nunit = int(math.floor(math.log(nbytes, 1024)))
    nsize = round(nbytes/(math.pow(1024, nunit)), 2)
    return f"{nsize:.2f} {unit[nunit]}"


def print_used_ram():
    ram = psutil.virtual_memory()
    print("CPU used memory:", format_bytes(ram[3]))
    print("CPU used memory (%):", f"{ram[3]*100/ram[0]:.2f}%")


def print_used_vram():
    nvmlInit()

    if nvmlDeviceGetCount() < 1:
        return None

    h = nvmlDeviceGetHandleByIndex(0)
    info = nvmlDeviceGetMemoryInfo(h)

    print(f"GPU used memory:", format_bytes(info.used))
    print(f"GPU used memory (%):", f"{info.used*100/info.total:.2f}%")


if __name__ == "__main__":
    print_used_ram()
    print_used_vram()