python - 优化 numpy 中的内存使用

Question

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

一张测试图像大约有 1000 万像素，灰度化后相当于 10MB。它被转换为 dtype uint8 的 Numpy 数组，经过一些处理(应用 Hamming 窗口)后，它是 dtype float64 的数组。两张图片以这种方式加载到数组中；后面的 FFT 步骤会生成 dtype complex128 的数组。在添加过多的 gc.collect 调用之前，程序内存大小往往会随着每一步而增加。此外，似乎大多数 Numpy 操作都会给出可用的最高精度的结果。

在我的 1GB Linux 机器上运行测试(没有 gc.collect 调用)导致长时间的颠簸，这是我没有等到的。我还没有详细的内存使用统计信息——我尝试了一些 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

This在 SO 上说“Scipy 0.8 将对几乎所有的 fft 代码提供单精度支持”，SciPy 0.8.0 beta 1 刚刚发布。
(自己没试过，胆小。)