Skip to content

Consider torch.float16 casting for cellpose segmentation workflows #29

Description

@lanery

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions