CODEWITHSUNDEEP
pythonnumpyprofilingfft
Zaccharie Ramzi
Zaccharie Ramzi

Reputation: 2316

numpy fftn very inefficient for 2d fft of several images

I wanted to compute the fourier transform of several images. I was therefore benchmarking numpy's fft.fftn against a brute force for loop.

This is the code I used to benchmark the 2 approaches (in a jupyter notebook):

import numpy as np

x = np.random.rand(32, 256, 256)

def iterate_fft(arr):
    k = np.empty_like(arr, dtype=np.complex64)
    for i, a in enumerate(arr):
        k[i] = np.fft.fft2(a)
    return k

k_it = iterate_fft(x)
k_np = np.fft.fftn(x, axes=(1, 2))
np.testing.assert_allclose(k_it.real, k_np.real)
np.testing.assert_allclose(k_it.imag, k_np.imag)

%%timeit
k_it = iterate_fft(x)

Output: 63.6 ms ± 1.23 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
k_np = np.fft.fftn(x, axes=(1, 2))

Output: 122 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Why is there such a huge difference ?

Upvotes: 3

Views: 604

Answers (2)

Zaccharie Ramzi
Zaccharie Ramzi

Reputation: 2316

So a person involved in the numpy fft development has answered the deep question on GitHub and it turns out that the slowdown is most likely coming from some multi dimensional array rearrangement used by pocketfft.

It will all be a memory when numpy switches to the scipy 1.4 implementation which can be shown using my benchmark to not have these drawbacks.

Upvotes: 2

Sam Mason
Sam Mason

Reputation: 16174

These routines in numpy seem to currently assume that the last dimension will always be the smallest. When this is actually true fftn is faster, sometimes by a lot.

That said, I get a much smaller difference in performance between these two methods than you (with Python 3.7.4, numpy 1.17.2). For your example, iterate_fft takes 46ms while ffn takes 50. But if I flip the axes around, to (256, 256, 32), I get 55ms and 40ms respectively. Pushing even further with a shape of (256, 256, 2) I get 21ms and 4ms respectively.

Note that if performance is really an issue, there are other FFT libraries available that perform better in some situations. Also the full fftpack in scipy can have very different performance than the more limited code in numpy.

Note that your usage of fftn basically does:

x = np.random.rand(32, 256, 256)

a = np.fft.fft(x, n=256, axis=2)
a = np.fft.fft(a, n=256, axis=1)

np.testing.assert_allclose(np.fft.fftn(x, axes=(1, 2)), a)

Upvotes: 1

Related Questions