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

Question

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

1)有没有简单的方法可以找到cudf dataframe (和其他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

这些值反映了使用的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)

大多数图形处理器操作都需要O(N)的暂存缓冲区，因此，如果最终的DataFrame或Series大于剩余可用内存，则可能会遇到RMM_OUT_OF_MEMORY错误。