v0.3.2-beta: Multi-GPU control, better reconstructed volumes, data inversion

This is a patch release to make some minor improvements based on using the cryoDRGN-AI software across a variety of different datasets since the initial public release, as well as feedback from other users.

Improving control over multi-GPU behavior

We have added the multigpu configuration parameter to the reconstruction model engine to allow the user to decide whether multiple GPUs should be used. Whereas previous versions of cryoDRGN-AI were only able to use all GPUs available, as of v0.3.2-b the default behavior is to use only one GPU (or one CPU if no GPUs are available) and activate multi-GPU training using e.g:

drgnai setup my_out_dir/ --particles particles.128.mrcs --ctf ctf.pkl --pose pose.pkl --reconstruction-type het --cfgs 'multigpu=True'
drgnai train my_out_dir/

We have also added the --multigpu flag to drgnai train to facilitate activating multi-GPU behavior:

drgnai setup my_out_dir/ --particles particles.128.mrcs --ctf ctf.pkl --pose pose.pkl --reconstruction-type het
drgnai train my_out_dir/ --multigpu

Producing better reconstructed volumes at each epoch

Reconstructed volumes my_out_dir/out/reconstruct.<epoch>.mrc at each epoch now correspond to the image closest to the centroid of the z-latent-space for heterogeneous reconstruction, whereas previously we simply used the first image in the dataset. We expect this to create a more representative snapshot of what is happening during model training, although we still expect users to apply analyze to better understand both the characteristics of the latent space and the full range of reconstructed volumes at the end of training.

Adding control over data inversion

Previous cryoDRGN-AI versions assumed inverted data, and depended on the invert configuration parameter to drgnai analyze to correct this in the case of datasets such as 50S (EMPIAR-10076) which are uninverted. We have now added the invert_data configuration parameter to the reconstruction engine that is set to True by default, but should be set to False in the case of uninverted input image stacks like 50S. This also allows the correct volumes to be reconstructed at each epoch in the case of 50S!

Other fixes and minor changes

• reducing default batch size for known poses to 32 from 16, and for SGD from 256 to 32, as the previous defaults tended to create out-of-memory issues for many users' hardware setups (see e.g. #11, #16)
• the parameters for data normalization — by default still calculated by sampling a subset of the input dataset — can now be set explicitly by the user with the new configuration parameters norm_mean and norm_std, which overrides the default calculation
• updating the dataset_fast module for new reconstruction engine APIs, thus fixing a bug when using fast_dataloading=True
• better log messages regarding multi-GPU usage, and updated batch sizes when using multiple GPUs:

  (INFO) (reconstruct.py) (30-Jun-25 13:19:12) Using one GPU in spite of 4 available GPUs because --multigpu is not being used!
  

  (INFO) (reconstruct.py) (30-Jun-25 13:19:12) Using 4 GPUs!
  (INFO) (reconstruct.py) (30-Jun-25 13:19:12) Increasing batch size for known poses to 64
  (INFO) (reconstruct.py) (30-Jun-25 13:19:12) Increasing batch size for HPS to 32
  (INFO) (reconstruct.py) (30-Jun-25 13:19:12) Increasing batch size for SGD to 128
  

• fixing total loss log messages at the end of pretraining to use the number of particles used in pretraining as the denominator when calculating average loss, as opposed to the total number of particles in the dataset as previously
• updating documentation for new package name: DRGN-AI -> cryoDRGN-AI