如何使用PyOpenCL檢查GPU內存是否可用

Question

我想知道是否有一種方法可以在函數使用前檢查多少GPU內存可用。 我的代碼經常使用1.5 GB或更多的GPU內存，如果我的程序想要使用GPU時其他東西正在使用GPU，則會收到MemoryError異常或類似的東西。

我想實現某種代碼，以便檢查GPU是否有足夠的可用內存，如果可以，請繼續運行，如果沒有，請等待直到可用。

（最好是，我想在嘗試使用GPU之前先進行檢查，而不是使用try-except循環，然后在失敗時重試）

我檢查了PyOpenCL文檔，以查看device_info下是否有相關的內容，但是找不到任何實際的描述。

Answer 1

這是不可能的，實際上是OpenCL的限制，而不僅僅是PyOpenCL。 看這里 。

在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()