rapids - 使用 rapids.ai 时如何确定内存统计信息？

Question

我正在使用 rapids.ai 的 python 库我开始怀疑的关键事情之一是:如何以编程方式检查内存分配？我知道我可以使用 nvidia-smi 查看一些总体高级统计数据，但我特别想知道:

1) 是否有一种简单的方法来查找 cudf 数据帧(和其他 rapids 对象)的内存占用？

2) 有没有办法确定可用的设备内存？

我确信 C++ 程序员有很多方法可以获取这些详细信息，但我希望找到一个能让我留在 Python 中的答案。

Answer 1

1) 用法

所有 cudf 对象都应该有 .memory_usage() 方法:

import cudf
x = cudf.DataFrame({'x': [1, 2, 3]})
x_usage = x.memory_usage(deep=True)
print(x_usage)

输出:

x        24
Index     0
dtype: int64

这些值反射(reflect)了使用的 GPU 内存。

2) 剩余

您可以使用 pynvml 读取剩余的可用 GPU 内存:

import pynvml
​
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # Need to specify GPU
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
mem.free, mem.used, mem.total
(33500299264, 557973504, 34058272768)

大多数 GPU 操作都需要 O(N) 的暂存缓冲区，因此如果您以 DataFrame 或 Series 结束，您可能会遇到 RMM_OUT_OF_MEMORY 错误 大于您剩余的可用内存。