CODEWITHSUNDEEP
pythoncupy
clemisch
clemisch

Reputation: 1034

How to fully release GPU memory used in function

I'm using cupy in a function that receives a numpy array, shoves it on the GPU, does some operations on it and returns a cp.asnumpy copy of it.

The problem: The memory is not freed after the function (as seen in ndidia-smi).

I know about the caching and re-using of memory done by cupy. However, this seems to work only per-user. When multiple users are computing on the same GPU-server, they are limited by the cached memory of other users.

I also tried calling cp._default_memory_pool.free_all_blocks() inside the function at the end. This seems to have no effect. Importing cupy in the main code and calling free_all_blocks "manually" works, but I'd like to encapsulate the GPU stuff in the function, not visible to the user.

Can you fully release GPU memory used inside a function so that it's usable by other users?

Minimal example:

Main module:

# 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()

Function module:

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

Upvotes: 7

Views: 10084

Answers (2)

Will
Will

Reputation: 369

You could use a decorator like this:

import cupy as cp

def free_gpu_memory(func):
    def wrapper_func(*args, **kwargs):
        retval = func(*args, **kwargs)
        cp._default_memory_pool.free_all_blocks()
        return retval
    return wrapper_func

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

All local variables defined in test_function will be out of scope when free_all_blocks() is called. This has the advantage that you don't need to keep track of which variables to delete in test_function

Upvotes: 3

kmaehashi
kmaehashi

Reputation: 1045

CuPy uses Python's reference counter to track which arrays are in use. In this case, you should del arr_gpu before calling free_all_blocks in test_function.

See here for more details: https://docs.cupy.dev/en/latest/user_guide/memory.html

Upvotes: 15

Related Questions