[BUG] Spike in memory usage end of cellfinder inference phase

[marcodallavecchia]

Describe the bug
At the end of cellfinder cell splitting phase, RAM usage spikes until it takes all my laptop RAM (at least 8Gb which is available) and the process crashes.

To Reproduce
I 100 slices of the demo data provided in the open software week.

Screenshots

Computer used (please complete the following information):

• OS: Ubuntu on Lenovo Thinkbook 16p Gen4
• Version 22.04
• Hardware specs 16 Gb RAM, Nvidia GPU GeForce RTX 4060 8Gb

[adamltyson]

@matham I don't suppose you've seen this? Following #440, it looks like the memory usage of structure splitting has increased. We didn't notice until we ran a course, where we saw that cellfinder is crashing on laptops as it now needs ~10-15GB of RAM for structure splitting.

[matham]

I have not seen it, but I'm running on systems with lots of memory. I imagine it has something to do with the parallelization of structure splitting, which happened around then. I'll take a look.

[matham]

Do you perchance have the data you tested it with? And the config values? I can use my own if not, but perhaps that especially uses more RAM.

[adamltyson]

It's possible that the issue is with classification, @IgorTatarnikov is going to dig into it a bit deeper.

The data was our tutorial data from the website, with default parameters, using the napari plugin. I think we must just not have tested it on laptops fully.

[marcodallavecchia]

I can confirm that running it through brainmapper with default parameters, I got the same with almost 12 Gb RAM free.

[IgorTatarnikov]

I compared v1.1.3 to v1.7.0 (I had to go all the way back to the tensorflow version to see a difference in behvaiour)

v1.7.0 - it used ~6-8 GB of RAM during candidate detection and cluster splitting with a spike up to ~20 GB of immediately after cluster splitting and before inference (stdout reported cluster splitting as complete but the torch inference progress bar hadn't showed up yet).

v1.1.3 - it used ~6 GB of RAM throughout the entire process from detection to inference. Might help pin down the source! I wonder if limiting the number of threads we use during candidate detection/splitting will have an effect on this spike?

[matham]

Actually, I think I did observe this. If we're talking about the inference part, on a CPU with 50+ cores, this exact thing happens.

Basically, in cellfinder, we specify the number of free CPU cores (

cellfinder/cellfinder/core/classify/classify.py

Line 115 in e97aec6

workers=workers,

). So if you have like 72 cores, PyTorch will try to use 70 cores. And loading the data and everything into all these subprocesses will exhaust resources.

That's why I prefer the model of specifying the maximum number of workers. This is in fact one of the things I changed in this PR #493 to use both the free and max cores to set the number of workers (https://github.com/matham/cellfinder/blob/bea7055b2e71b2990226f1f29ecfe841c278df0c/cellfinder/core/classify/classify.py#L97-L98).

FTR, I found with my PR, there wasn't much benefit to use more than like 6 workers for multiprocess data loading.

So I think this issue is probably with inference, not detection.

[alessandrofelder]

Next step:

• check on @marcodallavecchia with hard-coded parameter 6
• if fixed, expose the parameter.

[marcodallavecchia]

following up on this @adamltyson @alessandrofelder, indeed reducing the cores solves the issue.

I simply increased the number of free CPUs param from the Napari widget. I had to go all the way up to 16 free cores (using only 4 at a time on my 20 cores laptop CPU) to stay within the 16Gb of RAM. With 4 cores at a time I used a max of ~14Gb.

No need to call this a bug, but simply a thing to pay attention to. Do we want to make the number of cores specification more visible for different stages of the cellfinder pipeline, or do we leave it as it is?

[adamltyson]

I think we probably want to expose more parameters. E.g. we want to make sure that users can use many CPU cores for the detection step (that doesn't seem to use lots of memory), but mak sure the RAM usage during classification isn't so high.

[matham]

For detection phase, we do benefit from more cores, without much more resource usage if using CPU. If using GPU, it's like inference and doesn't benefit much beyond a few. Even on CPU, we don't benefit much beyond like 12 as you can see here
https://github.com/brainglobe/cellfinder/blob/main/cellfinder/core/detect/filters/setup_filters.py#L360. More than like 12 causes slowdowns due to thread contentions.

That's why I prefer the model to specify the (max) number of workers you want because that's typically more relevant. Which you can do indirectly via free cores. Also, detection typically can benefit from a bit more cores than inference, because inference is mostly purely GPU processing, with some light pre-processing.

[adamltyson]

We need to make some kind of change, because the current model isn't working (hence why this issue was raised in the first place). While most cellfinder analyses are carried out with high-core, high-memory systems with GPUs, most of the time a potential user initially encounters cellfinder is by testing it on a laptop, or in an environment they haven't set the GPU up on. We want to make sure cellfinder runs relatively quickly on these systems, but doesn't crash, so I'm not sure what the right way to set the default parameters is? Should we just explicitly set a max number of cores and make the default something like 6?

[alessandrofelder]

Should we just explicitly set a max number of cores and make the default something like 6?

I'd say so, with safeguards if the laptop has fewer than 6 cores