python - 如何使用 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()