内存分析#

峰值内存使用是 GPU 编程中的常见问题,因为 GPU 内存通常小于可用的 CPU 内存。为了方便识别内存热点,cuDF 提供了一个内存分析器。它会带来一定的开销,因此请避免在对性能敏感的代码中使用它。

启用内存分析#

首先,通过调用 rmm.statistics.enable_statistics() 在 RMM 中启用内存分析。这会向当前 RMM 内存资源添加一个统计信息资源适配器,从而使 cuDF 能够访问内存分析信息。有关更多详细信息,请参阅 RMM 文档

其次,通过将 memory_profiling 选项设置为 True 来在 cuDF 中启用内存分析。可以使用 cudf.set_option() 或在 Python 解释器启动前设置环境变量 CUDF_MEMORY_PROFILING=1

要获取分析结果,请使用 cudf.utils.performance_tracking.print_memory_report() 或通过使用 cudf.utils.performance_tracking.get_memory_records() 访问原始分析数据。

示例#

接下来,我们将启用分析,执行一些工作,然后打印分析结果

>>> import cudf
>>> from cudf.utils.performance_tracking import print_memory_report
>>> from rmm.statistics import enable_statistics
>>> enable_statistics()
>>> cudf.set_option("memory_profiling", True)
>>> cudf.DataFrame({"a": [1, 2, 3]})  # Some work
   a
0  1
1  2
2  3
>>> print_memory_report()  # Pretty print the result of the profiling
Memory Profiling
================

Legends:
ncalls       - number of times the function or code block was called
memory_peak  - peak memory allocated in function or code block (in bytes)
memory_total - total memory allocated in function or code block (in bytes)

Ordered by: memory_peak

ncalls memory_peak memory_total filename:lineno(function)
     1          32           32 cudf/core/dataframe.py:690(DataFrame.__init__)
     2           0            0 cudf/core/index.py:214(RangeIndex.__init__)
     6           0            0 cudf/core/index.py:424(RangeIndex.__len__)