更新时间:2024-09-05 GMT+08:00
分享

如何在代码中打印GPU使用信息

用户可通过shell命令或python命令查询GPU使用信息。

使用shell命令

  1. 执行nvidia-smi命令。

    依赖CUDA nvcc

    watch -n 1 nvidia-smi

  2. 执行gpustat命令。
    pip install gpustat 
    gpustat -cp -i

    使用Ctrl+C可以退出。

使用python命令

  1. 执行nvidia-ml-py3命令(常用)。
    !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)
        util = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
        mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
        print(f"|Device {i}| Mem Free: {mem.free/1024**2:5.2f}MB / {mem.total/1024**2:5.2f}MB | gpu-util: {util.gpu:3.1%} | gpu-mem: {util.memory:3.1%} |")

  2. 执行nvidia_smi + wapper + prettytable命令。

    用户可以将GPU信息显示操作看作一个装饰器,在模型训练过程中就可以实时的显示GPU状态信息。

    def gputil_decorator(func):
        def wrapper(*args, **kwargs):
            import nvidia_smi
            import prettytable as pt
    
            try:
                table = pt.PrettyTable(['Devices','Mem Free','GPU-util','GPU-mem'])
                nvidia_smi.nvmlInit()
                deviceCount = nvidia_smi.nvmlDeviceGetCount()
                for i in range(deviceCount):
                    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
                    res = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
                    mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
                    table.add_row([i, f"{mem.free/1024**2:5.2f}MB/{mem.total/1024**2:5.2f}MB", f"{res.gpu:3.1%}", f"{res.memory:3.1%}"])
    
            except nvidia_smi.NVMLError as error:
                print(error)
    
            print(table)
            return func(*args, **kwargs)
        return wrapper

  3. 执行pynvml命令。

    nvidia-ml-py3可以直接查询nvml c-lib库,而无需通过nvidia-smi。因此,这个模块比nvidia-smi周围的包装器快得多。

    from pynvml import *
    nvmlInit()
    handle = nvmlDeviceGetHandleByIndex(0)
    info = nvmlDeviceGetMemoryInfo(handle)
    print("Total memory:", info.total)
    print("Free memory:", info.free)
    print("Used memory:", info.used)

  4. 执行gputil命令。
    !pip install gputil
    import GPUtil as GPU
    GPU.showUtilization()

    import GPUtil as GPU
    GPUs = GPU.getGPUs()
    for gpu in GPUs:
        print("GPU RAM Free: {0:.0f}MB | Used: {1:.0f}MB | Util {2:3.0f}% | Total {3:.0f}MB".format(gpu.memoryFree, gpu.memoryUsed, gpu.memoryUtil*100, gpu.memoryTotal))

    注:用户在使用pytorch/tensorflow等深度学习框架时也可以使用框架自带的api进行查询。

相关文档