How to check the total memory consumed by a function in python?

Viewed 28

I want to check the total memory consumed by a function in python. I have used tracemalloc but using it you can only check current size and peak size of memory blocks. I would like to be able to measure the memory consumed like using the @btime macro in Julia, which returns the memory allocated before returning the value of the expression.

I tried also something like this:

# importing libraries
import os
import psutil
 
# inner psutil function
def process_memory():
    process = psutil.Process(os.getpid())
    mem_info = process.memory_info()
    return mem_info.rss
 
# decorator function
def profile(func):
    def wrapper(*args, **kwargs):
 
        mem_before = process_memory()
        result = func(*args, **kwargs)
        mem_after = process_memory()
        print("{}:consumed memory: {:,}".format(
            func.__name__,
            mem_before, mem_after, mem_after - mem_before))
 
        return result
    return wrapper
 
# instantiation of decorator function
@profile
 
# main code for which
# memory has to be monitored
def func():
    x = [1] * (10 ** 7)
    y = [2] * (4 * 10 ** 8)
    del x
    return y
 
func()

But this solution returns the consumed memory only at the beginning, and returns 0 on each subsequent call.

1 Answers

Take a look at the memory-profiler library.

Here an example taken from the docs:

# example.py

from memory_profiler import profile

@profile
def my_func():
    a = [1] * (10 ** 6)
    b = [2] * (2 * 10 ** 7)
    del b
    return a

if __name__ == '__main__':
    my_func()

Then in the command line type:

mprof run example.py
mprof plot

And you will get a table with stats and a nice looking plot: enter image description here

You should easily be able to adapt this example to your function of interest.

Related