sysinfo.py 文件源码

def getGPUUsage():
    try:
        pynvml.nvmlInit()
        count = pynvml.nvmlDeviceGetCount()
        if count == 0:
            return None

        result = {"driver": pynvml.nvmlSystemGetDriverVersion(),
                  "gpu_count": int(count)
                  }
        i = 0
        gpuData = []
        while i<count:
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
            gpuData.append({"device_num": i, "name": pynvml.nvmlDeviceGetName(handle), "total": round(float(mem.total)/1000000000, 2), "used": round(float(mem.used)/1000000000, 2)})
            i = i+1

        result["devices"] = jsonpickle.encode(gpuData, unpicklable=False)
    except Exception as e:
        result = {"driver": "No GPU!", "gpu_count": 0, "devices": []}

    return result