如何在 tensorflow 2.0b 中检查/释放 GPU 内存？答案

【问题标题】：How can I check/release GPU-memory in tensorflow 2.0b?如何在 tensorflow 2.0b 中检查/释放 GPU 内存？
【发布时间】：2019-12-05 18:24:53
【问题描述】：

在我的 tensorflow2.0b 程序中，我确实收到了这样的错误

    ResourceExhaustedError: OOM when allocating tensor with shape[727272703] and type int8 on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:TopKV2]

此程序中的许多基于 GPU 的操作已成功执行后，会发生此错误。

我喜欢释放与这些过去操作相关的所有 GPU 内存以避免上述错误。如何在 tensorflow-2.0b 中做到这一点？如何从我的程序中检查内存使用情况？

我只能使用 tf.session() 找到相关信息，这在 tensorflow2.0 中不再可用

【问题讨论】：

请看看这个issue 应该可以帮助您解决问题。我想如果你减少批量大小，应该注意这个错误。

标签： python-3.x gpu tensorflow2.0

【解决方案1】：

您可能有兴趣使用此Python 3 Bindings for the NVIDIA Management Library。

我会尝试这样的：

import nvidia_smi

nvidia_smi.nvmlInit()

handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
# card id 0 hardcoded here, there is also a call to get all available card ids, so we could iterate

info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)

print("Total memory:", info.total)
print("Free memory:", info.free)
print("Used memory:", info.used)

nvidia_smi.nvmlShutdown()

【讨论】：

import nvidia_smi 没有找到这个包。
此答案中链接的 github repo 提供了安装说明。使用 pip ： pip install nvidia-ml-py3 它安装和导入对我来说很好。