Transformación de Fourier 3D comparativamente lenta de Python numpy

Para mi trabajo, necesito realizar transformaciones de Fourier discretas (DFT) en imágenes grandes. En el ejemplo actual, necesito un FT 3D para una imagen de 1921 x 512 x 512 (junto con FFT 2D de 512 x 512 imágenes). En este momento, estoy usando el paquete numpy y la función asociada np.fft.fftn () . El fragmento de código a continuación muestra de manera ejemplar los tiempos FFT en 2D y 3D en una cuadrícula 2D / 3D generada por números aleatorios de igual tamaño / un poco más pequeño de la siguiente manera:

import sys import numpy as np import time tas = time.time() a = np.random.rand(512, 512) tab = time.time() b = np.random.rand(100, 512, 512) tbfa = time.time() fa = np.fft.fft2(a) tfafb = time.time() fb = np.fft.fftn(b) tfbe = time.time() print "initializing 512 x 512 grid:", tab - tas print "initializing 100 x 512 x 512 grid:", tbfa - tab print "2D FFT on 512 x 512 grid:", tfafb - tbfa print "3D FFT on 100 x 512 x 512 grid:", tfbe - tfafb 

Salida:

 initializing 512 x 512 grid: 0.00305700302124 initializing 100 x 512 x 512 grid: 0.301637887955 2D FFT on 512 x 512 grid: 0.0122730731964 3D FFT on 100 x 512 x 512 grid: 3.88418793678 

El problema que tengo es que necesitaré este proceso con bastante frecuencia, por lo que el tiempo empleado en cada imagen debería ser corto. Al realizar pruebas en mi propia computadora (computadora portátil de segmento medio, 2 GB de RAM asignada a una máquina virtual (-> por lo tanto, una cuadrícula de prueba más pequeña)), como puede ver, la FFT 3D toma ~ 5 s (orden de magnitud). Ahora, en el trabajo, las máquinas son mucho mejores, los sistemas de architecture de clúster / cuadrícula y FFT son mucho más rápidos. En ambos casos los 2D terminan casi instantáneamente.

Sin embargo, con 1921x512x512, np.fft.fftn () toma ~ 5 min. Como supongo que la implementación de scipy no es mucho más rápida y teniendo en cuenta que las FFT de MATLAB de cuadrículas del mismo tamaño finalizan en unos 5 s, mi pregunta es si existe un método para acelerar el proceso o casi hasta los tiempos de MATLAB. Mi conocimiento sobre las FFT es limitado, pero aparentemente MATLAB utiliza el algoritmo FFTW, que python no. ¿Alguna posibilidad razonable de que con algún paquete pyFFTW tenga tiempos similares? Además, 1921 parece una elección desafortunada, teniendo solo 2 factores primos (17, 113), por lo que asumo que esto también juega un papel importante. Por otro lado, 512 es una potencia adecuada de dos. ¿Se pueden lograr tiempos similares a los de MATLAB, si es posible, sin rellenar con ceros hasta 2048?

Lo pregunto porque tendré que usar FFT mucho (¡hasta una cantidad en la que tales diferencias serán de gran influencia!) Y en caso de que no haya posibilidad de reducir los tiempos de cálculo en Python, tendré que cambiar a otro , implementaciones más rápidas.

Sí, existe la posibilidad de que el uso de FFTW a través de la interfaz pyfftw reduzca el tiempo de cálculo en comparación con numpy.fft o scipy.fftpack . Los rendimientos de estas implementaciones de algoritmos DFT se pueden comparar en puntos de referencia como este : se reportan algunos resultados interesantes en Mejora del rendimiento de FFT en Python

Sugiero el siguiente código para una prueba:

 import pyfftw import numpy import time import scipy f = pyfftw.n_byte_align_empty((127,512,512),16, dtype='complex128') #f = pyfftw.empty_aligned((33,128,128), dtype='complex128', n=16) f[:] = numpy.random.randn(*f.shape) # first call requires more time for plan creation # by default, pyfftw use FFTW_MEASURE for the plan creation, which means that many 3D dft are computed so as to choose the fastest algorithm. fftf=pyfftw.interfaces.numpy_fft.fftn(f) #help(pyfftw.interfaces) tas = time.time() fftf=pyfftw.interfaces.numpy_fft.fftn(f) # here the plan is applied, nothing else. tas = time.time()-tas print "3D FFT, pyfftw:", tas f = pyfftw.n_byte_align_empty((127,512,512),16, dtype='complex128') #f = pyfftw.empty_aligned((33,128,128), dtype='complex128', n=16) f[:] = numpy.random.randn(*f.shape) tas = time.time() fftf=numpy.fft.fftn(f) tas = time.time()-tas print "3D FFT, numpy:", tas tas = time.time() fftf=scipy.fftpack.fftn(f) tas = time.time()-tas print "3D FFT, scipy/fftpack:", tas # first call requires more time for plan creation # by default, pyfftw use FFTW_MEASURE for the plan creation, which means that many 3D dft are computed so as to choose the fastest algorithm. f = pyfftw.n_byte_align_empty((128,512,512),16, dtype='complex128') fftf=pyfftw.interfaces.numpy_fft.fftn(f) tas = time.time() fftf=pyfftw.interfaces.numpy_fft.fftn(f) # here the plan is applied, nothing else. tas = time.time()-tas print "3D padded FFT, pyfftw:", tas 

Para un tamaño de 127 * 512 * 512, en mi computadora modesta, obtuve:

 3D FFT, pyfftw: 3.94130897522 3D FFT, numpy: 16.0487070084 3D FFT, scipy/fftpack: 19.001199007 3D padded FFT, pyfftw: 2.55221295357 

Entonces, pyfftw es significativamente más rápido que numpy.fft y scipy.fftpack . Usar el relleno es incluso más rápido, pero lo que se calcula es diferente.

Por último, pyfftw puede parecer más lento en la primera ejecución debido a que utiliza el indicador FFTW_MEASURE acuerdo con la documentación . Es bueno si, y solo si, se calculan sucesivamente muchos DFT del mismo tamaño.