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

我想知道在函数使用之前是否有办法检查可用的GPU内存量.我的代码通常使用1.5 GB或更多的GPU内存,如果我的程序想要使用它时其他东西正在使用GPU,我会得到一个MemoryError异常或类似的东西.

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

(最好,我想在尝试使用GPU之前检查,而不是使用try-except循环,如果失败则只重试)

我检查了PyOpenCL文档,看看在device_info下是否有相关内容,但我找不到任何实际的描述.

这是不可能的,实际上是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()
相关文章
相关标签/搜索