如何完全释放函数中使用的 GPU 内存

Question

我在接收numpy 数组的函数中使用cupy，将它推到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，然后再在test_function 中调用free_all_blocks。

更多详情请看这里： https://docs.cupy.dev/en/latest/user_guide/memory.html