cellpose icon indicating copy to clipboard operation
cellpose copied to clipboard

Add flag to compute masks on CPU but run network inference on GPU [FEATURE]

Open Cactusfacts opened this issue 1 year ago • 1 comments

Hi Cellpose team,

I've found that Cellpose can be very useful applied to large images (e.g. histology tissue). When the GPU version of the mask computation is run, this can result in a huge GPU memory spike and allocation failure as it attempts to move the entire flow field back onto the GPU. If I tile the input image to avoid this, I'll have mask artifacts at the seams.

Currently, Cellpose operation is either all-GPU or all-CPU. I think a finer-grained control to enable GPU network inference but CPU mask&flow steps would be a useful addition. I've added this as a flag in my local copy but wanted to share the idea in case it can help benefit others.

Thanks!

Cactusfacts avatar Apr 19 '23 16:04 Cactusfacts

thanks I think this is a good idea, cc @mrariden please add this to the to-do

carsen-stringer avatar May 09 '23 17:05 carsen-stringer