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内存是否可用

2 个答案: