Memory issues with Reference Based Motion Correction?

Motion Correction
3

Hi @larsonmattr and @rbs_sci. Sorry to hear you’re having VRAM problems with RBMC. We can try to address this at some point in the future but for now perhaps I can offer some ideas about reducing RBMC’s VRAM demands.

Very roughly, the baseline VRAM usage is proportional to:

  • Number of particles per movie
  • Number of frames per movie
  • Number of Fourier components in the alignment band. This warrants some explanation. The trajectory alignment is done in Fourier space, but not using all of the Fourier components (so the memory demand generally isn’t as bad as (box size)^2). In the hyperparameter estimation step, half of the Fourier components between 20A and the FSC cutoff frequency are used by default. In the two later steps, all of them between 20A and FSC frequency are used (which is why you might run out of VRAM on the dose weight stage even if hyperparam optimization succeeds).
  • Naturally, oversubscription doubles the required VRAM per GPU.

Perhaps counterintuitively, the number of particles used for dose weight computation and hyperparameter search don’t affect VRAM demands. 16-bit float also has no effect as all actual math is still done in float32. And the memory used for caching also affects only CPU RAM, not VRAM.

Some things that could be done to reduce these factors:

  • Turn off oversubscription

  • If you’re processing a heterogeneous dataset and you have a lot of particles per movie, you can reduce the effective number of particles per movie by processing each species in a separate reference motion job.

  • You can lower the “Fraction of FCs used for alignment” parameter to reduce the VRAM requirement of hyperparameter estimation.

  • You can turn off “Use all Fourier components” in the dose weight and final reconstruction stage parameters. This will cause those stages to use the same subset of the Fourier components that is used in hyperparameter estimation.

There’s also some VRAM set aside on one of the GPUs for slicing the volumes, and that’s proportional to the box size cubed. If your box has an excessive and unneeded amount of empty padding around it, cropping it may help with this.

3 Likes