Optimizar el uso de memoria en numpy

Question

El siguiente programa carga dos imágenes con PyGame, las convierte en matrices Numpy y luego realiza otras operaciones de Numpy (como FFT) para emitir un resultado final (de unos pocos números). Las entradas pueden ser grandes, pero en cualquier momento solo uno o dos objetos grandes deben estar en vivo.

Una imagen de prueba mide unos 10 millones de píxeles, lo que se traduce en 10 MB una vez que está en escala de grises. Se convierte a una matriz Numpy de dtype uint8, que después de algún procesamiento (aplicando ventanas de Hamming), es una matriz de dtype float64. Dos imágenes se cargan en matrices de esta manera; los pasos de FFT posteriores dan como resultado una matriz de tipo complex128. Antes de agregar las llamadas excesivas gc.collect, el tamaño de la memoria del programa tendía a aumentar con cada paso. Además, parece que la mayoría de las operaciones de Numpy darán un resultado con la mayor precisión disponible.

Ejecutando la prueba (sin las llamadas gc.collect) en mi máquina de 1GB de Linux da como resultado una manipulación prolongada, que no he esperado. Todavía no tengo estadísticas detalladas de uso de la memoria. Probé algunos módulos de Python y el comando time sin éxito; ahora estoy buscando valgrind. Ver PS (y lidiar con la falta de respuesta de la máquina en las últimas etapas de la prueba) sugiere un uso máximo de memoria de aproximadamente 800 MB.

Una matriz de 10 millones de celdas de complex128 debería ocupar 160 MB. Tener (idealmente) como mucho dos de estos productos al mismo tiempo, además de las bibliotecas no insustanciales de Python y Numpy y otros parafernalia, probablemente signifique permitir 500 MB.

puedo pensar en dos ángulos desde los que atacar el problema:

• descarte matrices intermedias tan pronto como sea posible. Para eso están las llamadas gc.collect: parecen haber mejorado la situación, ya que ahora se completa con solo unos minutos de aplastamiento ;-). Creo que uno puede esperar que la programación con mucha memoria en un lenguaje como Python requerirá alguna intervención manual.

• Uso de matrices Numpy con menos precisión en cada paso. Desafortunadamente, las operaciones que devuelven matrices, como fft2, no parecen permitir que se especifique el tipo.

Así que mi pregunta principal es: ¿Hay una manera de especificar la precisión en las operaciones de salida de matriz numpy?

Más en general, ¿hay otras técnicas comunes de conservación de la memoria cuando se utiliza Numpy?

Además, ¿tiene Numpy una forma más idiomática de liberar la memoria de la matriz? (Me imagino que esto dejaría el objeto de matriz en vivo en Python, pero en un estado inutilizable). La eliminación explícita seguida de GC inmediata se siente hacky.

import sys 
import numpy 
import pygame 
import gc 


def get_image_data(filename): 
    im = pygame.image.load(filename) 
    im2 = im.convert(8) 
    a = pygame.surfarray.array2d(im2) 
    hw1 = numpy.hamming(a.shape[0]) 
    hw2 = numpy.hamming(a.shape[1]) 
    a = a.transpose() 
    a = a*hw1 
    a = a.transpose() 
    a = a*hw2 
    return a 


def check(): 
    gc.collect() 
    print 'check' 


def main(args): 
    pygame.init() 

    pygame.sndarray.use_arraytype('numpy') 

    filename1 = args[1] 
    filename2 = args[2] 
    im1 = get_image_data(filename1) 
    im2 = get_image_data(filename2) 
    check() 
    out1 = numpy.fft.fft2(im1) 
    del im1 
    check() 
    out2 = numpy.fft.fft2(im2) 
    del im2 
    check() 
    out3 = out1.conjugate() * out2 
    del out1, out2 
    check() 
    correl = numpy.fft.ifft2(out3) 
    del out3 
    check() 
    maxs = correl.argmax() 
    maxpt = maxs % correl.shape[0], maxs/correl.shape[0] 
    print correl[maxpt], maxpt, (correl.shape[0] - maxpt[0], correl.shape[1] - maxpt[1]) 


if __name__ == '__main__': 
    args = sys.argv 
    exit(main(args)) 

Answer 1

This de eso dice "Scipy 0.8 tendrá soporte precisión simple para casi todo el código FFT", y SciPy 0.8.0 beta 1 es justo fuera.
(No lo he probado yo mismo, cobardemente.)

Answer 2

Si entiendo correctamente, está calculando una convolución entre dos imágenes. El paquete Scipy contiene un módulo dedicado para eso (ndimage), que podría ser más eficiente en cuanto a la memoria que el enfoque "manual" a través de las transformadas de Fourier. Sería bueno intentar usarlo en lugar de pasar por Numpy.