Numpy's bitwise operations and garbage collector

Question

I am working with numpy's bitwise functions and trying to reduce the memory footprint of my script:

import os, gc, psutil, resource, numpy as np

def print_memory():
    print(
        "Max %.2f MB - Current %.2f MB"
        % (
            resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
            psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024,
        )
    )

print_memory()

data = np.arange(1000000, dtype=np.uint64) # About 7.63 MB
result = np.zeros(data.size, dtype=np.uint64) # More 7.63 MB

a = np.right_shift(data, 50)
b = np.left_shift(data, 39)

result = np.bitwise_xor(data, np.bitwise_xor(np.bitwise_or(a, b), 1))

del a, b
gc.collect()

print_memory()

Which outputs:

Max 28.05 MB - Current 28.05 MB
Max 66.71 MB - Current 43.90 MB

This makes sense. At the end we have 43.9 MB − 28.05 MB = 15.85 MB, which is the size of my two 7.63 MB arrays. Same for the maximum memory used, 66.71 MB − 28.05 MB = 38.66 MB, which are my two 7.63 MB arrays (data and result), along with their copies saved in variables a and b, and the result given by np.bitwise_or(a, b). I.e., 7.63 MB * 5 = 38.15 MB, which is almost 38.66 MB.

The following changes and their results intrigue me.

import os, gc, psutil, resource, numpy as np

def print_memory():
    print(
        "Max %.2f MB - Current %.2f MB"
        % (
            resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
            psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024,
        )
    )

print_memory()

data = np.arange(1000000, dtype=np.uint64)  # About 7.63 MB
result = np.zeros(data.size, dtype=np.uint64)  # More 7.63 MB

result = np.bitwise_xor(
    data,
    np.bitwise_xor(np.bitwise_or(np.right_shift(data, 50), np.left_shift(data, 39)), 1),
)

gc.collect()

print_memory()

Which produces:

Max 28.14 MB - Current 28.14 MB
Max 58.98 MB - Current 51.52 MB

I assume that numpy/python was somehow able to reduce the maximum memory footprint from 5 times my input array to four, which is welcomed. However, why is it finishing with 51.52 MB − 28.14 MB = 23.38 MB, which is pretty much three arrays of size 7.63 MB? Where is the third 7.63 MB array coming from? Is it the result of an inner operation which is not being picked up by the garbage collector? Ultimately, how can I resolve this issue and have the best of both worlds: reduced maximum memory usage provided by the second script and "correct" memory consumption at the end as in the first script?

user_4542928 · Answer

Using tracemalloc you can verify that only two allocations remain after both executions. The problem is likely that 7MB is not enough for the OS or standard library to actually deallocate the memory, but free has been called.

See the code A:

import tracemalloc
import os, gc, psutil, resource, numpy as np

def print_memory():
    print(
        "Max %.2f MB - Current %.2f MB"
        % (
            resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
            psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024,
        )
    )


tracemalloc.start()
# ... start your application ...

snapshot1 = tracemalloc.take_snapshot()

print_memory()

data = np.arange(1000000, dtype=np.uint64) # About 7.63 MB
result = np.zeros(data.size, dtype=np.uint64) # More 7.63 MB

a = np.right_shift(data, 50)
b = np.left_shift(data, 39)

result = np.bitwise_xor(data, np.bitwise_xor(np.bitwise_or(a, b), 1))

del a, b
gc.collect()

print_memory()

snapshot2 = tracemalloc.take_snapshot()

top_stats = snapshot2.compare_to(snapshot1, 'lineno')

print("[ Top 10 differences ]")
for stat in top_stats[:10]:
    print(stat)

produces:

Max 27.09 MB - Current 27.09 MB
Max 65.30 MB - Current 42.52 MB
[ Top 10 differences ]
main.py:27: size=7813 KiB (+7813 KiB), count=4 (+4), average=1953 KiB
main.py:21: size=7813 KiB (+7813 KiB), count=3 (+3), average=2604 KiB
/usr/lib/python3/dist-packages/psutil/__init__.py:699: size=816 B (+816 B), count=3 (+3), average=272 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1390: size=720 B (+720 B), count=3 (+3), average=240 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1597: size=648 B (+648 B), count=2 (+2), average=324 B
/usr/lib/python3/dist-packages/psutil/__init__.py:365: size=576 B (+576 B), count=1 (+1), average=576 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:187: size=520 B (+520 B), count=6 (+6), average=87 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1596: size=456 B (+456 B), count=1 (+1), average=456 B
/usr/lib/python3/dist-packages/psutil/_common.py:337: size=424 B (+424 B), count=1 (+1), average=424 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1429: size=376 B (+376 B), count=4 (+4), average=94 B

As you can see only two allocations of 7MB remain.

See code B:

import tracemalloc
import os, gc, psutil, resource, numpy as np

def print_memory():
    print(
        "Max %.2f MB - Current %.2f MB"
        % (
            resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
            psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024,
        )
    )


tracemalloc.start()
# ... start your application ...

snapshot1 = tracemalloc.take_snapshot()

print_memory()

data = np.arange(1000000, dtype=np.uint64)  # About 7.63 MB
result = np.zeros(data.size, dtype=np.uint64)  # More 7.63 MB

result = np.bitwise_xor(
    data,
    np.bitwise_xor(np.bitwise_or(np.right_shift(data, 50), np.left_shift(data, 39)), 1),
)


gc.collect()

print_memory()

snapshot2 = tracemalloc.take_snapshot()

top_stats = snapshot2.compare_to(snapshot1, 'lineno')

print("[ Top 10 differences ]")
for stat in top_stats[:10]:
    print(stat)

This produces:

Max 27.08 MB - Current 27.08 MB
Max 57.55 MB - Current 50.14 MB
[ Top 10 differences ]
main2.py:26: size=7813 KiB (+7813 KiB), count=6 (+6), average=1302 KiB
main2.py:21: size=7813 KiB (+7813 KiB), count=3 (+3), average=2604 KiB
/usr/lib/python3/dist-packages/psutil/__init__.py:699: size=816 B (+816 B), count=3 (+3), average=272 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1390: size=720 B (+720 B), count=3 (+3), average=240 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1597: size=648 B (+648 B), count=2 (+2), average=324 B
/usr/lib/python3/dist-packages/psutil/__init__.py:365: size=576 B (+576 B), count=1 (+1), average=576 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:187: size=520 B (+520 B), count=6 (+6), average=87 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1596: size=456 B (+456 B), count=1 (+1), average=456 B
/usr/lib/python3/dist-packages/psutil/_common.py:337: size=424 B (+424 B), count=1 (+1), average=424 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1429: size=376 B (+376 B), count=4 (+4), average=94 B

Also only two allocations of 7MB remain.

A way to force the OS and stdlib to actually free memory is to use a large amount of it, so that release mechanisms trigger. For example for 76MB, you get the expected result:

code A, with data = np.arange(10000000, dtype=np.uint64) # About 76.3 MB gives:

Max 26.93 MB - Current 26.93 MB
Max 408.54 MB - Current 179.69 MB
[ Top 10 differences ]
main.py:27: size=76.3 MiB (+76.3 MiB), count=4 (+4), average=19.1 MiB
main.py:21: size=76.3 MiB (+76.3 MiB), count=3 (+3), average=25.4 MiB
/usr/lib/python3/dist-packages/psutil/__init__.py:699: size=816 B (+816 B), count=3 (+3), average=272 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1390: size=720 B (+720 B), count=3 (+3), average=240 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1597: size=648 B (+648 B), count=2 (+2), average=324 B
/usr/lib/python3/dist-packages/psutil/__init__.py:365: size=576 B (+576 B), count=1 (+1), average=576 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:187: size=520 B (+520 B), count=6 (+6), average=87 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1596: size=456 B (+456 B), count=1 (+1), average=456 B
/usr/lib/python3/dist-packages/psutil/_common.py:337: size=424 B (+424 B), count=1 (+1), average=424 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1429: size=376 B (+376 B), count=4 (+4), average=94 B

And code B with data = np.arange(10000000, dtype=np.uint64) # About 76.3 MB gives:

Max 27.13 MB - Current 27.13 MB
Max 332.43 MB - Current 179.89 MB
[ Top 10 differences ]
main2.py:26: size=76.3 MiB (+76.3 MiB), count=6 (+6), average=12.7 MiB
main2.py:21: size=76.3 MiB (+76.3 MiB), count=3 (+3), average=25.4 MiB
/usr/lib/python3/dist-packages/psutil/__init__.py:699: size=816 B (+816 B), count=3 (+3), average=272 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1390: size=720 B (+720 B), count=3 (+3), average=240 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1597: size=648 B (+648 B), count=2 (+2), average=324 B
/usr/lib/python3/dist-packages/psutil/__init__.py:365: size=576 B (+576 B), count=1 (+1), average=576 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:187: size=520 B (+520 B), count=6 (+6), average=87 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1596: size=456 B (+456 B), count=1 (+1), average=456 B
/usr/lib/python3/dist-packages/psutil/_common.py:337: size=424 B (+424 B), count=1 (+1), average=424 B
/usr/lib/python3/dist-packages/psutil/_pslinux.py:1429: size=376 B (+376 B), count=4 (+4), average=94 B

user_10229754 · Answer

On my machine, your two methods give the following footprints:

Method 1:

print_memory()

data = np.arange(1000000, dtype=np.uint64) # About 7.63 MB
result = np.zeros(data.size, dtype=np.uint64) # More 7.63 MB

a = np.right_shift(data, 50)
b = np.left_shift(data, 39)

result = np.bitwise_xor(data, np.bitwise_xor(np.bitwise_or(a, b), 1))

del a, b
gc.collect()

print_memory()

This produces:

Max 29.96 MB - Current 29.96 MB
Max 68.26 MB - Current 45.40 MB

Method 2:

print_memory()

data = np.arange(1000000, dtype=np.uint64)  # About 7.63 MB
result = np.zeros(data.size, dtype=np.uint64)  # More 7.63 MB

result = np.bitwise_xor(
    data,
    np.bitwise_xor(np.bitwise_or(np.right_shift(data, 50), np.left_shift(data, 39)), 1),
)

gc.collect()

print_memory()

This produces:

Max 29.88 MB - Current 29.88 MB
Max 60.46 MB - Current 52.95 MB

Method 3:

Now for a potential solution. It seems like it's possible to get the best of both worlds if the operations are written the following way:

print_memory()

data = np.arange(1000000, dtype=np.uint64)  # About 7.63 MB
result = np.zeros(data.size, dtype=np.uint64)  # More 7.63 MB

np.bitwise_or(np.right_shift(data, 50), np.left_shift(data, 39), out = result)
np.bitwise_xor(result, 1, out = result)
np.bitwise_xor(data, result, out = result)

gc.collect()

print_memory()

This produces:

Max 29.96 MB - Current 29.96 MB
Max 60.41 MB - Current 45.41 MB

Now I'm not absolutely sure why this works. However, I noticed that these bitwise operations support the out argument; I wondered if setting out = result would help numpy reuse the same memory location in each step, and thus increase memory efficiency somehow. Seems like it helps.

Numpy's bitwise operations and garbage collector

2 Answers