Как полностью освободить память графического процессора, используемую в функции ⇐ Python

[Anonymous]

Я использую cupy в функции, которая получает массив numpy, помещает его в графический процессор, выполняет с ним некоторые операции и возвращает его копию cp.asnumpy.

Проблема: память не освобождается после выполнения функции (как видно из ndidia-smi).

Я знаю о кэшировании и повторном использовании памяти, выполняемом cupy. Однако, похоже, это работает только для каждого пользователя. Когда несколько пользователей выполняют вычисления на одном и том же GPU-сервере, они ограничены кэш-памятью других пользователей.

Я также пробовал вызвать cp._default_memory_pool.free_all_blocks() внутри функции в конце. Кажется, это не имеет никакого эффекта. Импорт cupy в основной код и вызов free_all_blocks «вручную» работает, но я бы хотел инкапсулировать данные графического процессора в функцию, невидимую для пользователя.

Можете ли вы полностью освободить память графического процессора, используемую внутри функции, чтобы ее могли использовать другие пользователи?



Минимальный пример:

Основной модуль:

# 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


Подробнее здесь: https://stackoverflow.com/questions/535 ... n-function