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

Question

我在接收 numpy 数组的函数中使用 cupy，将其推送到 GPU 上，对其进行一些操作，然后 returns 一个 cp.asnumpy它的副本。

问题：函数执行后内存没有释放（见ndidia-smi）。

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

我也试过最后在函数内部调用 cp._default_memory_pool.free_all_blocks()。这似乎没有效果。在主代码中导入 cupy 并调用 free_all_blocks "manually" 可以，但我想将 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 的引用计数器来跟踪正在使用的数组。在这种情况下，您应该在 test_function.

中调用 free_all_blocks 之前 del arr_gpu

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

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

How to fully release GPU memory used in function

python

cupy