Как сделать 100000 раз 2d FFT быстрее с помощью Python? - PullRequest
Новая
       16

Как сделать 100000 раз 2d FFT быстрее с помощью Python?

1 голос
/ 06 марта 2019

У меня есть 3d-массив с массивностью формы (100000, 256, 256), и я хотел бы сделать FFT для каждого стека массива 2d, что означает 100000 раз FFT.

Я проверил скорость одиночных и суммированных данных с минимальным кодом ниже. 

import numpy as np
a = np.random.random((256, 256))
b = np.random.random((10, 256, 256))

%timeit np.fft.fft2(a)

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

Что дает следующее:

872 µs ± 19.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

6.46 ms ± 227 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

100000 раз fft займет более одной минуты.

Есть ли более быстрый способ сделать несколько FFT или IFFT одновременно?

Обновление: После небольшого поиска я нашел cupy , который, кажется, может помочь.

1 Ответ

1 голос
/ 06 марта 2019

pyfftw , обтекание библиотеки FFTW , вероятно, быстрее, чем библиотека FFTPACK , обернутая np.fft и scipy.fftpack , В конце концов, FFTW означает быстрое преобразование Фурье на Западе.

Минимальный код: 

import numpy as np
import pyfftw
import multiprocessing
b = np.random.random((100, 256, 256))
bb = pyfftw.empty_aligned((100,256, 256), dtype='float64')
bf= pyfftw.empty_aligned((100,256, 129), dtype='complex128')
fft_object_b = pyfftw.FFTW(bb, bf,axes=(1,2),flags=('FFTW_MEASURE',), direction='FFTW_FORWARD',threads=multiprocessing.cpu_count())
bb=b
fft_object_b(bb)

Вот расширенный код, определяющий время выполнения np.fft и pyfftw: 

import numpy as np
from timeit import default_timer as timer
import multiprocessing
a = np.random.random((256, 256))
b = np.random.random((100, 256, 256))

start = timer()
for i in range(10):
    np.fft.fft2(a)
end = timer()
print"np.fft.fft2, 1 slice", (end - start)/10

start = timer()
for i in range(10):
     bf=np.fft.fftn(b, axes=(1, 2,))
end = timer()
print "np.fft.fftn, 100 slices", (end - start)/10
print "bf[3,42,42]",bf[3,42,42]


import pyfftw

aa = pyfftw.empty_aligned((256, 256), dtype='float64')
af= pyfftw.empty_aligned((256, 129), dtype='complex128')
bb = pyfftw.empty_aligned((100,256, 256), dtype='float64')
bf= pyfftw.empty_aligned((100,256, 129), dtype='complex128')
print 'number of threads:' , multiprocessing.cpu_count()

fft_object_a = pyfftw.FFTW(aa, af,axes=(0,1), flags=('FFTW_MEASURE',), direction='FFTW_FORWARD',threads=multiprocessing.cpu_count())

fft_object_b = pyfftw.FFTW(bb, bf,axes=(1,2),flags=('FFTW_MEASURE',), direction='FFTW_FORWARD',threads=multiprocessing.cpu_count())


aa=a
bb=b
start = timer()
for i in range(10):
    fft_object_a(aa)
end = timer()
print "pyfftw, 1 slice",(end - start)/10

start = timer()
for i in range(10):
    fft_object_b(bb)
end = timer()
print "pyfftw, 100 slices", (end - start)/10
print "bf[3,42,42]",bf[3,42,42]

Наконец, результатом является значительное ускорение: pyfftw работает в 10 раз быстрее, чем np.fft на моем компьютере. , используя 2 потока. 

np.fft.fft2, 1 slice 0.00459032058716
np.fft.fftn, 100 slices 0.478203487396
bf[3,42,42] (-38.190256258791734+43.03902512127183j)
number of threads: 2
pyfftw, 1 slice 0.000421094894409
pyfftw, 100 slices 0.0439268112183
bf[3,42,42] (-38.19025625879178+43.03902512127183j)

Ваш компьютер кажется намного лучше моего!

...