Rendimiento de PyFFTW en matrices multidimensionales

Question

Tengo una matriz nD, digamos de dimensiones: (144, 522720) y necesito calcular su FFT.

PyFFTWparece más lento que numpyy scipy, que NO se espera.

¿Estoy haciendo algo obviamente mal?

A continuación está mi código

import numpy
import scipy      
import pyfftw
import time

n1 = 144
n2 = 522720
loops = 2

pyfftw.config.NUM_THREADS = 4
pyfftw.config.PLANNER_EFFORT = 'FFTW_ESTIMATE'
# pyfftw.config.PLANNER_EFFORT = 'FFTW_MEASURE'

Q_1 = pyfftw.empty_aligned([n1, n2], dtype='float64')
Q_2 = pyfftw.empty_aligned([n1, n2], dtype='complex_')
Q_ref = pyfftw.empty_aligned([n1, n2], dtype='complex_')

# repeat a few times to see if pyfft planner helps
for i in range(0,loops):
    Q_1 = numpy.random.rand(n1,n2)

    s1 = time.time()
    Q_ref = numpy.fft.fft(Q_1, axis=0)
    print('NUMPY - elapsed time: ', time.time() - s1, 's.')

    s1 = time.time()
    Q_2 = scipy.fft.fft(Q_1, axis=0)
    print('SCIPY - elapsed time: ', time.time() - s1, 's.')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

    s1 = time.time()
    Q_2 = pyfftw.interfaces.numpy_fft.fft(Q_1, axis=0)
    print('PYFFTW NUMPY - elapsed time = ', time.time() - s1, 's.')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

    s1 = time.time()
    Q_2 = pyfftw.interfaces.scipy_fftpack.fft(Q_1, axis=0)
    print('PYFFTW SCIPY - elapsed time = ', time.time() - s1, 's.')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

    s1 = time.time()
    fft_object = pyfftw.builders.fft(Q_1, axis=0)
    Q_2 = fft_object()
    print('FFTW PURE Elapsed time = ', time.time() - s1, 's')
    print('Equal = ', numpy.allclose(Q_2, Q_ref))

Answer 1

En primer lugar, si enciende la caché antes del bucle principal, las interfaces funcionan en gran medida como se esperaba:

pyfftw.interfaces.cache.enable()
pyfftw.interfaces.cache.set_keepalive_time(30)

Es interesante que a pesar de la sabiduría que debería almacenarse, la construcción de los pyfftwobjetos sigue siendo bastante lenta cuando el caché está desactivado. No importa, este es exactamente el propósito del caché. En su caso, debe hacer que el tiempo de mantenimiento de la caché sea bastante largo porque su ciclo es muy largo.

En segundo lugar, no es una comparación justa incluir el tiempo de construcción del fft_objecten la prueba final. Si lo mueve fuera del temporizador, entonces llamar fft_objectes una mejor medida.

En tercer lugar, también es interesante ver que incluso con el caché activado, la llamada a numpy_fftes más lenta que la llamada a scipy_fft. Dado que no hay una diferencia obvia en la ruta del código, sugiero que es un problema de almacenamiento en caché. Este es el tipo de problema que timeitbusca mitigar. Aquí está mi código de tiempo propuesto que es más significativo:

import numpy
import scipy
import pyfftw
import timeit

n1 = 144
n2 = 522720

pyfftw.config.NUM_THREADS = 4
pyfftw.config.PLANNER_EFFORT = 'FFTW_MEASURE'

Q_1 = pyfftw.empty_aligned([n1, n2], dtype='float64')

pyfftw.interfaces.cache.enable()
pyfftw.interfaces.cache.set_keepalive_time(30)

times = timeit.repeat(lambda: numpy.fft.fft(Q_1, axis=0), repeat=5, number=1)
print('NUMPY fastest time = ', min(times))

times = timeit.repeat(lambda: scipy.fft.fft(Q_1, axis=0), repeat=5, number=1)
print('SCIPY fastest time = ', min(times))

times = timeit.repeat(
    lambda: pyfftw.interfaces.numpy_fft.fft(Q_1, axis=0), repeat=5, number=1)
print('PYFFTW NUMPY fastest time = ', min(times))

times = timeit.repeat(
    lambda: pyfftw.interfaces.scipy_fftpack.fft(Q_1, axis=0), repeat=5, number=1)
print('PYFFTW SCIPY fastest time = ', min(times))

fft_object = pyfftw.builders.fft(Q_1, axis=0)
times = timeit.repeat(lambda: fft_object(Q_1), repeat=5, number=1)
print('FFTW PURE fastest time = ', min(times))

En mi máquina, esto da una salida como:

NUMPY fastest time =  0.6622681759763509
SCIPY fastest time =  0.6572431400418282
PYFFTW NUMPY fastest time =  0.4003451430471614
PYFFTW SCIPY fastest time =  0.40362057799939066
FFTW PURE fastest time =  0.324020683998242

Puede hacerlo un poco mejor si no lo fuerza a copiar la matriz de entrada en un tipo de datos complejo cambiando Q_1a complex128:

NUMPY fastest time =  0.6483533839927986
SCIPY fastest time =  0.847397351055406
PYFFTW NUMPY fastest time =  0.3237176960101351
PYFFTW SCIPY fastest time =  0.3199474769644439
FFTW PURE fastest time =  0.2546963169006631

Esa interesante scipyralentización es repetible.

Dicho esto, si su entrada es real, debería estar haciendo una transformación real (para> 50% de aceleración pyfftw) y manipulando la salida compleja resultante.

Lo interesante de este ejemplo es (creo) lo importante que es la caché en los resultados (lo que sugiero es por qué cambiar a una transformación real es tan efectivo para acelerar las cosas). También ve algo dramático cuando cambia el tamaño de la matriz a 524288 (la siguiente potencia de dos, que cree que quizás podría acelerar las cosas, pero no ralentizarlas dramáticamente). En este caso, todo se ralentiza bastante, sobre scipytodo. Me parece que scipyes más sensible a la caché, lo que explicaría la ralentización al cambiar la entrada a complex128(522720 es un número bastante bueno para FFTing, por lo que quizás deberíamos esperar una desaceleración).

Finalmente, si la velocidad es secundaria a la precisión, siempre puede usar flotantes de 32 bits como tipo de datos. Si combina eso con hacer una transformación real, obtiene una aceleración mejor que el factor 10 sobre la numpymejor inicial dada anteriormente:

PYFFTW NUMPY fastest time =  0.09026529802940786
PYFFTW SCIPY fastest time =  0.1701313250232488
FFTW PURE fastest time =  0.06202622700948268

(numpy y scipy no cambian mucho, ya que creo que usan flotantes de 64 bits internamente).

Editar: Olvidé que las fftpackFFT reales de Scipy tienen una estructura de salida extraña, que se pyfftwreplica con cierta desaceleración. Esto se cambia para que sea más sensato en el nuevo módulo FFT .

La nueva interfaz FFT está implementada en pyFFTW y debería preferirse. Desafortunadamente, hubo un problema con la reconstrucción de los documentos, por lo que los documentos estaban desactualizados durante mucho tiempo y no mostraban la nueva interfaz; con suerte, eso se ha solucionado ahora.