如何充分释放函数中使用的GPU内存

Question

我cupy在一个接收numpy数组的函数中使用，将它推到 GPU 上，对其执行一些操作并返回它的cp.asnumpy副本。

问题：函数执行后内存未释放（如 中所示ndidia-smi）。

我知道cupy. 但是，这似乎仅适用于每个用户。当多个用户在同一个 GPU 服务器上进行计算时，他们会受到其他用户缓存内存的限制。

最后我还尝试cp._default_memory_pool.free_all_blocks()在函数内部调用。这似乎没有效果。导入cupy主代码并free_all_blocks“手动”调用有效，但我想将 GPU 内容封装在函数中，对用户不可见。

您能否完全释放函数内部使用的 GPU 内存，以便其他用户可以使用它？

---

最小的例子：

主要模块：

# dont import cupy here, only numpy
import numpy as np

# module in which cupy is imported and used
from memory_test_module import test_function

# host array
arr = np.arange(1000000)

# out is also on host, gpu stuff happens in test_function
out = test_function(arr)

# GPU memory is not released here, unless manually:
import cupy as cp
cp._default_memory_pool.free_all_blocks()

功能模块：

import cupy as cp

def test_function(arr):
    arr_gpu = cp.array(arr)
    arr_gpu += 1
    out_host = cp.asnumpy(arr_gpu)

    # this has no effect
    cp._default_memory_pool.free_all_blocks()

    return out_host

Answer 1

CuPy 使用 Python 的引用计数器来跟踪正在使用的数组。在这种情况下，你应该del arr_gpu收到主叫free_all_blocks在test_function。

有关更多详细信息，请参见此处：https : //docs.cupy.dev/en/latest/user_guide/memory.html