mxnet cause gpu memory leak inference with dynamic shape input, while using flask

Question

Sorry, I don't put on any code to prove

The problem is:

when I use a simple for loop like this:

for _ in range(2000):
    rnum = random.randint(1, 5)
    img = np.random.rand(rnum, 3, 112, 112)
    mxnet_model.inference(img)

it will work fine

However, if I cover code above with a flask API

It will cause gpu memory leak... which is very terrible

Answer 1

The "leak" you discovered is due to an inappropriate memory pooling strategy for this use case. Every time a new space is released, to avoid frequent calls to memory allocation, MXNet tries to cache that space in the memory pool for use next time. The default memory pooling strategy in MXNet is to use an exact size match, which is not suitable for the case of changing shapes due to high cache miss. You can switch to the round strategy by setting environment variable MXNET_GPU_MEM_POOL_TYPE=Round.

See reference: https://mxnet.apache.org/versions/1.6/api/faq/env_var#memory-options

Answer 2

Thanks for Szha's answer.

But setting up the environment variable MXNET_GPU_MEM_POOL_TYPE doesn't actually work for me.

I end up filling the input to a fixed size and peeling off the random extra outputs.