Perfomance questions regarding gpu saturation and use of num_cpus parameter

[bHimes]

Problem

When testing LEM on a single rtx4090 I cannot find a combination of num_cpus and orientation_batch_size that quite saturates the GPU based on gpu utilization as observed using nvidia-smi

Best config so far

• orientation_batch_size=48
• num_cpu=[ this doesn't seem to have an effect ]

I'm using the prebuilt version installed using pip in a python3.10 venv as suggested at the homepage under the pre-package releases.

• As an aside, it would be cool if the example code blocks were copyable. There are many mechanisms to do this, e.g.

Questions

1. What is the orientation_batch_size you typically recommend?
2. Is this just sending a stack of images off in a batch to pyTorch? Does that make the answer in the last question effectively "As many as will fit in memory."
3. Is num_cpus actually doing anything right now? It doesn't seem to be based on either my hardware resource monitors, or what I can make out from the code.

• I see that the num_cpus is not being validated as the GPU is here
• I see that you test that num_cpus is at least configurable and shouldn't allow negative numbers, but don't see any test on usage otherwise.

Any tips or tricks on how to get the most out of a given hardware setup would be very helpful. Thanks!

[mgiammar]

Hi Ben, thanks for pointing some of these things out! We plan on modifying some of the internals to better scale across different hardware setups including multi-node HPC environments, but right now we are focused on getting v1.0 released as the first stable version of the package.

Question Answers

1. We'd currently recommend the largest orientation_batch_size that's possible to fit in memory. For a 48GB A6000 ada with a 4k by 4k image, this is an orientation batch size of 8 (when using 13 defocus planes).
2. Yes, this is just sending off a stack of projections in PyTorch with batch shape orientation_batch_size * num_defocus. Ideally, we want to decouple the defocus from the batch size of things computed under-the-hood.
3. The num_cpus parameter currently does nothing, but this probably will change to mean the number of individual threads/streams submitting work a GPU in the future.

Configuration setup

Best config so far
orientation_batch_size=48

I'm curious what your problem setup looks like, that is:

• Image size
• Template shape
• Defocus search parameters

Notes on performance

Right now, the PyTorch backend works well for very large images on GPUs with lots of VRAM mostly because that's what we have available to the lab. Smaller images, like the extracted boxes in refine_template will not perform as well just because of PyTorch overhead on single-threaded programs. There is also a lot of memory overhead...

We are also hoping to move away from the pure PyTorch backend for match_template to really maximize hardware usage, e.g. using FastFFT for the cross-correlations ;)

Copyable code blocks

I will push some new docs later today including the copy button, thanks!

[mgiammar]

@bHimes I have a brief update on the num_cpus parameter and backend cross-correlation behavior which stems from some minor testing I did with cross-correlating using multiple streams vs multiple batches in Pytorch. We found that executing individual 2D cross-correlations across N streams was faster than doing a single, batched cross-correlation with N elements. Updating the match_template backend to use streams brings our computation time on an RTX A6000 ada from 11 hours to 8 hours.

The num_cpus parameter now controls the number of streams on each GPU over which cross-correlations are calculated. This also comes with a significant peak memory reduction which may help saturate your RTX 4090 cards.

See #56 for changes to the codebase.