优化numpy中的内存使用

下面的程序用PyGame加载两个图像,将它们转换为Numpy数组,然后执行一些其他Numpy操作(如FFT)以发出(几个数字的)最终结果。输入可以很大,但在任何时候都只能有一个或两个大对象处于活动状态。

一个测试图像大约是10万像素,一旦灰度化,它就会转换成10兆字节。它被转换成一个数据类型的Numpy数组 uint8 ,它经过一些处理(应用汉明窗口),是一个数据类型数组 float64 complex128 . 在添加多余的 gc.collect 调用时,程序内存大小往往随着每一步的增加而增加。此外,似乎大多数Numpy操作都会以最高的精度给出结果。

运行测试(sans) 调用)在我的1GB Linux机器上导致长时间的颠簸,这是我没有等待的。我还没有详细的内存使用统计数据——我尝试了一些Python模块和 time

一个1000万个单元的complex128数组应该占用160MB。拥有(理想情况下)最多两个这样的库同时存在,再加上并非虚无缥缈的Python和Numpy库以及其他工具,可能意味着允许500mb。

我可以从两个角度来解决这个问题:

• 尽快丢弃中间数组。这就是 gc.收集

• 在每一步使用精度较低的Numpy数组。不幸的是,返回数组的操作 fft2

有没有办法在Numpy数组操作中指定输出精度?

一般来说,在使用Numpy时,还有其他常见的内存保存技术吗?

另外,Numpy是否有更惯用的释放数组内存的方法(我想这会使array对象活在Python中,但处于不可用的状态。)显式删除之后立即执行GC感觉很不好。

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))