cellpose
cellpose copied to clipboard
Add flag to compute masks on CPU but run network inference on GPU [FEATURE]
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!
thanks I think this is a good idea, cc @mrariden please add this to the to-do