Eating up all your memory in Python? How to use ChatGPT to help out.

Recently, I was working on a project where I needed to optimize a hefty real-time AI edge compute task. The iterations were running around 15 seconds each at steady state. After some optimizations, the task was running in less than a second.

But... after about 150 iterations, the program slowed to a crawl. 15, 30, 60, 95 seconds per iteration. After using the jtops utility (a fantastic tool if you use Nvidia Jetson devices), it was pretty clear that my Python memory management was bad.

So I went to ChatGPT for help. Here's how you can too by giving the correct context and questions.

Set up profiling

Really simple here - import the python memory profiler:

from memory_profiler import profile

and then tag on the profiler decorator to the functions you are interested in:

@profile
def process_frame(
    stop_processes: mp.Value = None,
    queue_of_interesting_data: mp.Queue = None,
    model_weights_path: str = None,
    function_to_make_a_pretty_graph: Callable[[float], None] = None,
    data_input_fpath: Path = None,
):
...code...
return

Execute the script

Then run your script (on a jetson nano if you are still living in the late 20-teens like me)

(.venv) jetson@nano:~/dev/$ python3.8 scripts/my_nn_script.py  --weights_path models/my_model.tar --input_frame_path ../rt_data/ > profile.log

Once your script is completed, the end of profile.,log will spit out something that looks like this:

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
   333    278.1 MiB    278.1 MiB           1   @profile
   334                                         def process_frame(
   335                                             stop_processes: mp.Value = None,
   336                                             queue_of_interesting_data: mp.Queue = None,
   337                                             model_weights_path: str = None,
   338                                             function_to_make_a_pretty_graph: Callable[[float], None] = None,
   339                                             data_input_fpath: Path = None,
   340                                         )
   
...
   369                                                 # TODO: optimize this code
   370                                                 # create a simple map of uuids to the frames contained therein
   371   2287.5 MiB  -3613.1 MiB          50           check_map = {}
   372   2287.5 MiB -19290.5 MiB         204           for id in identities.values():
   373   2287.5 MiB -102576.4 MiB         919               check_map[id.uuid] = set([obs.frame.uuid for obs in id.observations])
   374   2287.5 MiB -15677.2 MiB         154               if id.uuid not in stale_identity_map:
   375   2287.5 MiB   -658.0 MiB           9                   stale_identity_map[str(id.uuid)] = [
   376   2287.5 MiB  -2701.0 MiB          38                       set([obs.frame.uuid for obs in id.observations]),
   377   2287.5 MiB   -658.0 MiB           9                       0,
   378                                                         ]
   ...
   421   2112.6 MiB   -168.2 MiB           1       return

Ask AI about It

Go to chat.openai.com (honestly, after a few hours with it, bard.google.com doesn't seem worth your time at this point. Less accurate and thorough. Probably a blog for another day though. Forgive me, google SEO overlords...)

The most effective question I found after some trials was:

how do I optimize this python code given this profiling output from using the memory_profile python library? <paste profile.log output for a function here>

The response that I got was:

Based on the profiling output, it seems that the main bottleneck in the process_frame function is in lines 372-373, where a dictionary is being created with UUIDs as keys and sets of UUIDs as values. This operation is causing a significant increase in memory usage.

To optimize this code, one possible approach is to use a defaultdict instead of a regular dictionary to avoid checking if the key already exists in the dictionary. Another approach is to use a set comprehension instead of a for loop to create the set of UUIDs for each identity.

Here's an example of how you could modify lines 372-373 to use a defaultdict and a set comprehension:

and then it dropped this code sample:

import collections

# create a defaultdict with sets as default values
check_map = collections.defaultdict(set)

# use a set comprehension to create the set of UUIDs for each identity
for id in identities.values():
    check_map[id.uuid].update(obs.frame.uuid for obs in id.observations)

That's undoubtedly a good suggestion to save memory! Additionally, the response is without context on the functions being called underneath the hood. If I give more information to ChatGPT, I would expect a dramatically improved suggestion about my data structures and implementations.

Hope this helps!

If you need help moving large ML/AI models to the real time edge, shoot us an email at info@depotanalytics.co for a free consultation. We'd love to help!