AdvancedPhotonSource / tike

Repository for ptychography software
http://tike.readthedocs.io
Other
29 stars 15 forks source link

Reduce memory footprint of Lamino and USFFT #68

Closed carterbox closed 4 years ago

carterbox commented 4 years ago

Purpose

This pull request addresses #66. Basically, the memory footprint of the Lamino operator is too large.

Approach

There are two approaches used in this pull request.

First, the checkerboard method for FFT shifting is implemented. This method uses properties of the FFT such that multiplying an array with a checkerboard of -1, 1 before and after the FFT is applied is equivalent to performing an FFT shift before and after.

Second, I dropped in the CuPy FFT functions that allow for the overwrite_x parameter which enables in-place FFT operations if available.

These two approaches have reduced the memory consumption of the Lamino operator by approximately half. Here is the memory usage before the changes:

tests/operators/test_lamino.py:49:test_adjoint (9.37GB, 5.57GB)
tests/operators/test_lamino.py:51:test_adjoint (8.06GB, 686.00MB)

Here is the memory usage after the changes:

tests/operators/test_lamino.py:49:test_adjoint (4.47GB, 3.42GB)
tests/operators/test_lamino.py:51:test_adjoint (3.87GB, 0.00B)

This is for n=16*14, ntheta=8*14, tilt=np.pi/3, eps=1e-3 with the Lamino adjoint test.

Pre-Merge Checklists

Submitter

Reviewer

nikitinvv commented 4 years ago

@carterbox , looks good. Thanks!