MouseLand/cellpose

[BUG] Batch size is ignored in 3D Cellpose 3D

Opened this issue · 0 comments

Hi, great work on the new v3 Cellpose. However I have a new issue with batching:

I could be missing something, but no matter what batch size I set, my GPU memory usage stays fixed at ~3.8GB. I'm operating on a Windows 11 workstation with an NVIDIA RTX A5500 (24GB) so this is really not making good use of the GPU. CUDA is enabled and the GPU is indeed running (checked via Task manager). The same code previously worked to change GPU memory usage depending on batch size in Cellpose v2.

This:

`masks, flows, styles, diams = model.eval(img, channels=[1,2], batch_size=240, diameter=60, do_3D=True, tile = True, cellprob_threshold = -1,
                                             anisotropy= anis_factor, min_size= 600, resample=False, stitch_threshold = 0.15, normalize = True, interp=False)`

Gives the same GPU usage and times as:

masks, flows, styles, diams = model.eval(img, channels=[1,2], batch_size=8, diameter=60, do_3D=True, tile = True, cellprob_threshold = -1,
                                             anisotropy= anis_factor, min_size= 600, resample=False, stitch_threshold = 0.15, normalize = True, interp=False)

torch version: 2.1.2