rapids - 在使用 rapids.ai 时如何确定内存统计数据？

Question

我正在使用rapids.ai的python 库，我开始想知道的关键问题之一是：如何以编程方式检查内存分配？我知道我可以nvidia-smi用来查看一些总体高级统计数据，但具体来说，我想知道：

1）有没有一种简单的方法可以找到 cudf 数据帧（和其他急流对象）的内存占用？

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

这些值反映了使用的 GPU 内存。

2) 剩余

您可以使用以下命令读取剩余的可用 GPU 内存pynvml：

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) 的暂存缓冲区，因此RMM_OUT_OF_MEMORY如果最终得到DataFrames 或Series大于剩余可用内存，则可能会遇到错误。