如何使用 PyOpenCL 检查 GPU 内存是否可用

Question

我想知道是否有一种方法可以在函数使用之前检查有多少 GPU 内存可用。我的代码经常使用 1.5 GB 或更多的 GPU 内存，如果当我的程序想要使用 GPU 时其他东西正在使用它，我会得到一个 MemoryError 异常或类似的东西。

我想实现某种代码，以便我可以检查 GPU 是否有足够的可用内存，如果有，请继续运行，但如果没有，请等到它可用。

（最好是，我想在尝试使用 GPU 之前进行检查，而不是使用 try-except 循环并在失败时重试）

我查看了 PyOpenCL 文档以查看 device_info 下是否有相关内容，但我找不到任何实际描述。

Answer 1

这是不可能的，实际上是 OpenCL 的限制，而不仅仅是 PyOpenCL。参见 here。

在 NVIDIA 设备上，您可以使用 nvidia-ml-py。然后你可以这样做：

from pynvml import *
nvmlInit()
for i in range(nvmlDeviceGetCount()):
    handle = nvmlDeviceGetHandleByIndex(i)
    meminfo = nvmlDeviceGetMemoryInfo(handle)
    print("%s: %0.1f MB free, %0.1f MB used, %0.1f MB total" % (
        nvmlDeviceGetName(handle),
        meminfo.free/1024.**2, meminfo.used/1024.**2, meminfo.total/1024.**2))
nvmlShutdown()

Answer 2

我尝试将上述答案与 python3 一起使用，但没有成功，但是我可以使用以下代码解决这个问题。

from py3nvml.py3nvml import *
nvmlInit()
deviceCount = nvmlDeviceGetCount()
for i in range(deviceCount):
    handle = nvmlDeviceGetHandleByIndex(i)
    gpuUtilization = nvmlDeviceGetUtilizationRates(handle)
    print(f'{nvmlDeviceGetName(handle)} Utilization: {gpuUtilization.gpu}% Memory: {gpuUtilization.memory}%')
nvmlShutdown()

如何使用 PyOpenCL 检查 GPU 内存是否可用

How to check if GPU memory is available using PyOpenCL

python

pyopencl