优化numpy中的内存使用

Question

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

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

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

complex128 的 1000 万个元胞数组应占用 160 MB。 （理想情况下）最多同时拥有其中两个，加上非实质性的 Python 和 Numpy 库和其他用具，可能意味着允许 500 MB。

我能想到两个角度来解决这个问题：

• 尽快丢弃中间数组。这就是gc.collect 呼吁的目的——他们似乎已经改善了这种情况，因为它现在只需几分钟的颠簸即可完成;-)。我认为人们可以预期，使用 Python 这样的语言进行内存密集型编程将需要一些手动干预。

• 在每个步骤中使用不太精确的 Numpy 数组。不幸的是，返回数组的操作（如 fft2）似乎不允许指定类型。

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

更一般地说，在使用 Numpy 时还有其他常见的节省内存的技术吗？

此外，Numpy 是否有更惯用的释放数组内存的方式？ （我想这会使数组对象在 Python 中存活，但处于不可用状态。）显式删除后立即 GC 感觉很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 因此，“Scipy 0.8将有几乎所有FFT代码的单一精度支持”， SciPy 0.8.0 beta 1 刚刚发布。
（自己没试过，胆小。）

Answer 2

如果我理解正确，您正在计算两个图像之间的卷积。 SCIPY软件包包含一个专用模块（ndimage），它可能比通过傅里叶变换的“手动”方法更多的内存效率。尝试使用它而不是经过numpy会很好。