Как полностью освободить память GPU, используемую в функции

Question

Я использую cupy в функции, которая получает numpy массив, пихает его на графический процессор, делает некоторые операции с ним и возвращает 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.chainer.org/en/stable/upgrade.html