python - cuda 功能的自动内存管理

Question

我一直在尝试使用一些用于 python 的 cuda 库来进行线性代数运算。迄今为止最直接的一个似乎是cupy。

然而，我遇到的所有这些问题是，每当操作变得足够大以对 gpu 计算感兴趣时，我就会遇到内存错误。

Cupy 或类似的库中是否有某种方法可以自动进行内存管理？

标签： pythonmemory-managementcupy

Answer 1

您可以尝试让 CuPy 在进入内存密集型区域之前回收所有未使用的 GPU 内存：

mempool = cupy.get_default_memory_pool()
mempool.free_all_blocks()

任何在调用之前引用计数为零的 CuPy 数组都将被释放，并将底层 GPU 内存返回到内存池。