Summary
Consider adding support for casting image data to torch.float16 (half precision) for cellpose segmentation workflows to reduce memory usage.
Motivation
Cellpose models can consume significant GPU memory when processing large images or batches. Using half-precision floats could substantially reduce memory footprint while maintaining acceptable segmentation quality.
Proposal
Add an option (at minimum) to cast input images to torch.float16 before passing to cellpose models. This could be:
- A parameter in segmentation functions (e.g.,
use_fp16=False)
- A global configuration option
- Automatic detection based on available GPU memory
Advantages
- Reduced memory usage: ~50% reduction in GPU memory for image tensors
- Larger batch sizes: Process more images simultaneously
- Faster inference: Potential speedups on modern GPUs with tensor cores
- Handle bigger images: Work with larger microscopy images that might otherwise OOM
Disadvantages
- Numerical precision: May impact segmentation accuracy for edge cases
- Compatibility: Not all GPUs support float16 efficiently (mainly older hardware)
- Additional complexity: Requires careful dtype management and testing
- Cellpose compatibility: Need to verify cellpose models work correctly with float16 inputs
Implementation considerations
- Should be opt-in initially to avoid breaking existing workflows
- May need validation/benchmarking to ensure segmentation quality is maintained
- Consider automatic fallback to float32 if float16 causes issues
Summary
Consider adding support for casting image data to
torch.float16(half precision) for cellpose segmentation workflows to reduce memory usage.Motivation
Cellpose models can consume significant GPU memory when processing large images or batches. Using half-precision floats could substantially reduce memory footprint while maintaining acceptable segmentation quality.
Proposal
Add an option (at minimum) to cast input images to
torch.float16before passing to cellpose models. This could be:use_fp16=False)Advantages
Disadvantages
Implementation considerations