python查看显存占用情况以及使用numba.cuda释放显存

一、python查看显存占用情况

# pip install pynvml 
def nvidia_info():
    nvidia_dict = {
        "state": True,
        "nvidia_version": "",
        "nvidia_count": 0,
        "gpus": []
    }
    try:
        nvmlInit()
        nvidia_dict["nvidia_version"] = nvmlSystemGetDriverVersion()
        nvidia_dict["nvidia_count"] = nvmlDeviceGetCount()
        for i in range(nvidia_dict["nvidia_count"]):
            handle = nvmlDeviceGetHandleByIndex(i)
            memory_info = nvmlDeviceGetMemoryInfo(handle)
            gpu = {
                "gpu_name": nvmlDeviceGetName(handle),
                "total": memory_info.total,
                "free": memory_info.free,
                "used": memory_info.used,
                "temperature": f"{nvmlDeviceGetTemperature(handle, 0)}℃",
                "powerStatus": nvmlDeviceGetPowerState(handle)
            }
            nvidia_dict['gpus'].append(gpu)
    except NVMLError as _:
        nvidia_dict["state"] = False
    except Exception as _:
        nvidia_dict["state"] = False
    finally:
        try:
            nvmlShutdown()
        except:
            pass
    return nvidia_dict



def check_gpu_mem_usedRate():
    info = nvidia_info()
    # print(info)
    used = info['gpus'][0]['used']
    tot = info['gpus'][0]['total']
    print(f"GPU0 used: {used}, tot: {tot}, 使用率：{used/tot}")

二、使用numba.cuda释放显存

释放显卡缓存

from  numba import cuda
device = cuda.get_current_device()
device.reset()

from  numba import cuda
cuda.select_device(1)
cuda.close()

参考 ： https://www.codenong.com/39758094/