use batches in make templates to not overload GPU memory

[jamesmkrieger]

At the moment, this is a limiting step for finer sampling

[jamesmkrieger]

This is the error. I get it with viewing distance 1.0 / (4.0 * pi) on my laptop and something similar using 0.5 / (4.0 * pi) on my workstation.

 Traceback (most recent call last):
  File "/home/jkrieger/software/scipion3/plugins/scipion-em-reweighting/reweighting/protocols/scripts/cryolike_template.py", line 26, in <module>
    make_templates_from_inputs(
  File "/home/jkrieger/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/make_templates_from_inputs_api.py", line 235, in make_templates_from_inputs
    (tp, name) = _make_raw_template(input, i, descriptor, t_float, device, verbose)
  File "/home/jkrieger/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/make_templates_from_inputs_api.py", line 168, in _make_raw_template
    tp = _make_templates_from_mrc_file(input, descriptor, t_float, device, verbose)
  File "/home/jkrieger/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/make_templates_from_inputs_api.py", line 77, in _make_templates_from_mrc_file
    return Templates.generate_from_physical_volume(
  File "/home/jkrieger/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/template.py", line 382, in generate_from_physical_volume
    templates_fourier = volume_phys_to_fourier_points(
  File "/home/jkrieger/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/microscopy/nufft.py", line 279, in volume_phys_to_fourier_points
    image_polar = torch.zeros((n_images, n_points), dtype = torch_complex_type, device = output_device)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.36 GiB. GPU 0 has a total capacity of 6.00 GiB of which 383.62 MiB is free. Including non-PyTorch memory, this process has 17179869184.00 GiB memory in use. Of the allocated memory 4.38 GiB is allocated by PyTorch, and 5.05 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

On the workstation, it says
Tried to allocate 11.26 GiB. GPU 0 has a total capacity of 10.75 GiB of which 10.48 GiB is free.

[jamesmkrieger]

Here's the full command and error log, running in iPython. I think the issue is that the GPU memory is already allocated in a previous step.

(cryolike-0.0.1) jkrieger@CNB-WHIPPLE ~/ScipionUserData/projects/TestPilar_copy $ ipython
Python 3.10.16 (main, Dec 11 2024, 16:24:50) [GCC 11.2.0]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.30.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import os
   ...: import numpy as np
   ...:
   ...: from cryolike.stacks import make_templates_from_inputs
   ...:
   ...: # Assume you have already run set_image_parameters_example with the same
   ...: # value of folder_output.
   ...: folder_output = 'Runs/015674_ReweightingProtComputeLikelihood/extra/templates/'
   ...: image_parameters_filename = os.path.join(folder_output, "parameters.npz")
   ...: list_of_inputs = ['Runs/001021_XmippProtConvertPdb/extra/7bnm.mrc']
   ...: verbose = True

In [2]: make_templates_from_inputs(
   ...:     list_of_inputs = list_of_inputs,
   ...:     image_parameters_file=image_parameters_filename,
   ...:     folder_output = folder_output,
   ...:     verbose = verbose
   ...: )
CUDA is available, using GPU.
Processing Runs/001021_XmippProtConvertPdb/extra/7bnm.mrc...
mrc_name: 7bnm
CUDA is available, using GPU.
---------------------------------------------------------------------------
OutOfMemoryError                          Traceback (most recent call last)
Cell In[2], line 1
----> 1 make_templates_from_inputs(
      2     list_of_inputs = list_of_inputs,
      3     image_parameters_file=image_parameters_filename,
      4     folder_output = folder_output,
      5     verbose = verbose
      6 )

File ~/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/make_templates_from_inputs_api.py:235, in make_templates_from_inputs(list_of_inputs, image_parameters_file, output_plots, folder_output, verbose, use_cuda, i_start)
    233 template_file_list = []
    234 for i, input in enumerate(list_of_inputs, start=i_start):
--> 235     (tp, name) = _make_raw_template(input, i, descriptor, t_float, device, verbose)
    236     tp.normalize_images_fourier(ord=2, use_max=False)
    237     plotter_fn(tp, descriptor, name)

File ~/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/make_templates_from_inputs_api.py:168, in _make_raw_template(input, iteration_cnt, descriptor, t_float, device, verbose)
    166 if extension in MRC_EXTENSIONS:
    167     print("mrc_name:", name)
--> 168     tp = _make_templates_from_mrc_file(input, descriptor, t_float, device, verbose)
    169 elif extension in PDB_EXTENSIONS:
    170     print(f"pdb_name: {name}")

File ~/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/make_templates_from_inputs_api.py:77, in _make_templates_from_mrc_file(mrc_file, descriptor, torch_float_type, device, verbose)
     75     raise ValueError(f"Can't happen: parsing mrc file {mrc_file} did not generate a physical density.")
     76 volume.density_physical = volume.density_physical.to(torch_float_type).to(device)
---> 77 return Templates.generate_from_physical_volume(
     78     volume=volume,
     79     polar_grid=descriptor.polar_grid,
     80     viewing_angles=descriptor.viewing_angles,
     81     precision=descriptor.precision,
     82     verbose=verbose,
     83     use_cuda=use_cuda
     84 )

File ~/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/stacks/template.py:382, in Templates.generate_from_physical_volume(cls, volume, polar_grid, viewing_angles, precision, use_cuda, nufft_eps, verbose)
    380 volume.density_physical = volume.density_physical.to(device)
    381 fourier_slices = _get_fourier_slices(polar_grid, viewing_angles, float_type=torch_float_type, device=device)
--> 382 templates_fourier = volume_phys_to_fourier_points(
    383     volume = volume,
    384     fourier_slices = fourier_slices,
    385     eps = nufft_eps,
    386     precision = Precision.SINGLE if precision == Precision.DEFAULT else precision,
    387     use_cuda = True,
    388     output_device = device,
    389     verbose = verbose
    390 )
    392 origin = torch.tensor([0.0, 0.0, 0.0], dtype = torch_float_type, device = device).unsqueeze(0)
    393 centers = volume_phys_to_fourier_points(
    394     volume = volume,
    395     fourier_slices = origin,
   (...)
    400     verbose = verbose
    401 )

File ~/software/scipion3/software/em/reweighting-0.0.1/cryolike/src/cryolike/microscopy/nufft.py:279, in volume_phys_to_fourier_points(volume, fourier_slices, isign, eps, precision, use_cuda, output_device, verbose)
    277 density = volume.density_physical.to(torch_complex_type) * rescale_factor
    278 if use_cuda:
--> 279     image_polar = torch.zeros((n_images, n_points), dtype = torch_complex_type, device = output_device)
    280     from cufinufft import Plan
    281     # from cufinufft import nufft3d2

OutOfMemoryError: CUDA out of memory. Tried to allocate 1.36 GiB. GPU 0 has a total capacity of 6.00 GiB of which 383.62 MiB is free. Including non-PyTorch memory, this process has 17179869184.00 GiB memory in use. Of the allocated memory 4.38 GiB is allocated by PyTorch, and 5.05 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)