使用 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

这些值反映了使用的 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) 的暂存缓冲区，因此如果您以 DataFrames 或 Series 大于您剩余的可用内存。

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

How do you determine memory stats while using rapids.ai?

rapids

cudf

1) 用法

2) 剩余