py4dstem/py4DSTEM

GPU Out of memory

ZOUCHEN158 opened this issue · 4 comments

Describe the bug
Out of memory allocating 8,589,934,592 bytes (allocated so far: 55,834,706,944 bytes)

py4DSTEM version
0.14.14
Python version
3.10
Operating system
Windows

GPU
image

  • Cupy Version
    image

Screenshots
image
image

Additional context
When I use ptychography's multislice reconstruction, an out of memory error is shown. I'm not sure if it's because the NVIDIA4090 I'm using is not suitable or for some other reason? Do I have to use a computing card like the A100? I installed tensorflow, but when I checked it didn't have access to the GPU. So I don't know what is the reason for failure?
image

This means that you do not have enough GPU memory to fit the entire reconstruction. There are a few ways to limit the memory usage:

  • Set storage='cpu' when creating the ptychography object, which will enable streaming to the GPU.
  • Set max_batch_size to be smaller in reconstruct.
  • Ensure that other programs are not allocating GPU memory. It appears from your screenshots that you have another Python interpreter also using the GPU. In order to see detailed memory usage information you may have to run nvidia-smi from an Administrator shell.

I don’t seem to see the relevant definition of storage=cpu. My py4dstem version is 0.14.8, and there is no definition of max_batch_size.

Hey @ZOUCHEN158

Indeed, the storage flag (and the associated memory tricks @sezelt recommends above) were introduced in version 0.14.9. I suggest upgrading to the latest version if you can: storing the amplitudes on the CPU and transferring only the necessary batch to the GPU for calculations has enabled very large FOVs at a modest performance cost.

I am hoping to write a tutorial notebook on memory management very soon.

Closing this issue, now that the tutorials have been updated to include memory management. Feel free to re-open if the issue persists.