在modelarts的notebook中如何在代码中打印gpu使用信息？-九游平台

依赖cuda nvcc

watch -n 1 nvidia-smi

pip install gpustat 

gpustat -cp -i

使用ctrl c可以退出。

!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%} |")

用户可以将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

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)

!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进行查询。