Runtime differences in version 0.1.3 and 0.1.4. [related to batch_size=None and training on large data - see issue #365]

[YSqkr-ep]

...
Hi,

I've noticed some pretty drastic runtime differences between cell2location versions 0.1.3 and 0.1.4. I ran the two versions on the same dataset with the same parameters. A process that took ~4hrs with 0.1.3 is taking ~40hrs with 0.1.4. Each iteration is also much longer with 0.1.4. (14s/it vs. 1s/it) One thing I noticed is that when i look at GPU usage with 'nvidia-smi', 0.1.3 is using ~5000MiB while the process run with 0.1.4. is only taking ~700MiB.
Are there any suggestions to make runtime better? I'm puzzled as to why the runtimes are so different between the versions. I've included some images of the runs.
Additionally, the "use_gpu" parameter in version 0.1.3. was pretty handy and it would be great to see it back in future versions.
Thanks!

[vitkl]

Hi @Qkrdudtj-ep

Do you mean the difference between version 0.1.3 and the latest version 0.1.5 or the previous version 0.1.4?

accellerator="gpu" has the same effect and is consistent with PyTorch Lightning and scvi-tools - so it makes sense to keep it.

This issue might be related to mismatches in torch version, cuda version and drivers on the host. We noticed that some GPU devices (H100) had issues and needed a different set of versions.

In general, we don't recommend training cell2location in minibatches because it reduces accuracy. See this discussion #356 and the tutorial note about full data training. I have never seen minibatch run times less than many days and mostly > 1 week. It's possible that the 0.1.3 version that you have has a bug that ignores batch_size and runs training on full data - run time for training on full data should be a few hours like you see in the second screenshot.

[vitkl]

We should probably make setting batch_size to anything other than None much harder.

[YSqkr-ep]

Hi Thank you for the response.
I was referring to the differences between version 0.1.3. and 0.1.4.
The version of torch and cuda i have with cell2location 0.1.4 are as follows.
Torch version: 2.5.1+cu124
CUDA (built with): 12.4
Device name: NVIDIA GeForce RTX 2080 Ti
cuDNN version: 90100

The reason I set a batch_size is because I run into an 'out of memory error' when i try setting batch_size to None.

With 'use_gpu', I can set which GPU I want to run on but I can't seem to do this with accelerator="gpu" and device=1. Is there a way around this?

[vitkl]

I would check scvi-tools and Lightning documentation and forum