Expose pin_memory to main and napari

[matham]

Description

This depends on #542.

What is this PR

• Bug fix
• Addition of a new feature
• Other

Why is this PR needed?

pin_memory is already used in detection, and will be used in classification. It should be exposed so users can use it to speed up GPU data uploading.

I didn't initially add it to classification as well. But I added it because we will need it and we will need to add it to brainglobe-workflows to pass it on. And making workflows being able to pass this parameter to both detection and classification, even before we start using it, will reduce compatibility issues.

What does this PR do?

It exposes it in napari and in the main functions. But it defaults to not being enabled due to the requirement of the data memory not being paged out of RAM if it's used. So users will need to turn it ON manually.

References

#540

How has this PR been tested?

Checked that it looks ok in napari.

Is this a breaking change?

No. It does change the default for detection from pinning to not pinning. But I added the default when we moved detection to pytorch.

Does this PR require an update to the documentation?

No.

Checklist:

• The code has been tested locally
• Tests have been added to cover all new functionality (unit & integration)
• The documentation has been updated to reflect any changes
• The code has been formatted with pre-commit

[alessandrofelder]

Thanks, @matham - the changes look straightforward and sensible to me, in principle.

Would you be able to help me understand

• what would happen if we enable pin_memory and the data gets swapped out of RAM?
  • does this just slow down things, or would this crash, or would this silently fail and give a wrong result?
• what would be a good way of knowing beforehand for users whether to switch pin_memory on or off (if there is one)?

[matham]

To the best of my understanding on how this works (this is a good overview). When we set pin_memory to True, it prevents the system from swapping out that memory.

So if you try to create a pin_memoryarray larger than available pure RAM, you'd probably get a memory error because we can't use paged memory for it. If you have an array with pin_memory taking up a lot of RAM and you need more RAM for other purposes, more than is available in pure RAM, the pin_memory RAM can't be swapped out so instead the new memory would have to be paged a lot, slowing down the system.

So it effectively reduces your system RAM by the amount used by pin_memory. This is helpful if you have enough RAM, but the system may page out the loaded arrays for whatever reasons, slowing down GPU transactions. But if you have a ton of RAM the system probably won't page it out anyway so it won't help a lot.

E.g. on our system, with 768GB RAM, running classification on 27,578,752 candidate cells took 33721s without pinning and 33401s with pinning (with #493). Which is about 5 min speedup with pinning memory. But I expect with a system with much less RAM and perhaps slower SSDs so paging is slower, not turning ON pin_memory would slow down classification more since the arrays may be paged out.

Side note, you don't want to pin memory in the main thread, as that may block, slowing down the pipeline. And we don't do it for detection or classification, unless you specifically set classification to not use any workers.

[matham]

With a better batch size, classifying 27578643 cells took 22092s (6.14 hours) with pinning memory and 24397s (6.78 hours) without. A 38 min difference. Which is not insignificant.

[alessandrofelder]

Thanks for the detailed explanation.
Confirming I get an 99%+ match on internal large test dataset MS_cx_left so will merge - thanks @matham !