py4dstem / py4DSTEM

GNU General Public License v3.0
199 stars 135 forks source link

GPU Out of memory #638

Closed ZOUCHEN158 closed 4 months ago

ZOUCHEN158 commented 5 months ago

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

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

sezelt commented 5 months ago

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:

ZOUCHEN158 commented 5 months ago

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.

gvarnavi commented 5 months ago

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.

gvarnavi commented 4 months ago

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