在numpy中优化内存使用

Question

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

测试图像大约为10M像素,一旦灰度化就会转换为10MB.它被转换为dtype的Numpy数组uint8,经过一些处理(应用Hamming窗口),是一个dtype数组float64.以这种方式将两个图像加载到阵列中; 稍后的FFT步骤会产生一个dtype数组complex128.在添加过多gc.collect调用之前,程序存储器大小趋于随着每一步增加.此外,似乎大多数Numpy操作都会以最高的精度给出结果.

gc.collect在我的1GB Linux机器上运行测试(没有调用)导致长时间的颠簸,我没有等待.我还没有详细的内存使用统计信息 - 我尝试了一些Python模块和time命令无济于事; 现在我正在研究valgrind.观看PS(以及在测试的后期阶段处理机器无响应)表明最大内存使用量约为800 MB.

一个1000万单元的complex128阵列应该占用160 MB.(理想情况下)最多有两个这样的活动,加上非实质性的Python和Numpy库以及其他随身用品,可能意味着允许500 MB.

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

• 尽快丢弃中间阵列.这就是所gc.collect要求的 - 他们似乎已经改善了这种情况,因为它现在只用几分钟的捶打就完成了;-).我认为可以预期像Python这样的语言中的内存密集型编程需要一些人工干预.

• 在每一步使用不太精确的Numpy数组.不幸的是,返回数组的操作fft2似乎不允许指定类型.

所以我的主要问题是:有没有办法在Numpy数组操作中指定输出精度？

更一般地说,使用Numpy时还有其他常见的内存保存技术吗？

另外,Numpy有更多惯用的释放阵列内存的方法吗？(我想这会让数组对象在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))

Answer 1

这 上面说“Scipy 0.8 将为几乎所有 fft 代码提供单精度支持”，并且 SciPy 0.8.0 beta 1 刚刚发布。
（自己没试过，胆怯。）