如何避免使用numpy vectorize时出现巨大的额外内存消耗？

Question

下面的代码最能说明我的问题：

控制台输出显示（注意，甚至第一个测试运行也需要大约8分钟），512x512x512的16位数组分配的内存使用量没有超过预期（每个大约256MB），而且通过查看“top”命令，进程的内存使用量通常保持在600MB以下，这也是预期的。

但是，当调用向量化版本的函数时，进程的内存使用量却变得非常大（超过7GB！）。即使我能想到的最明显的解释——向量化可能在内部将输入和输出转换为float64——也只能解释几GB的内存使用，尽管向量化函数返回的是int16，返回的数组肯定也是int16。有没有什么方法可以避免这种情况？我是不是对vectorize的otypes参数理解错了？

import numpy as np
import subprocess

def logmem():
    subprocess.call('cat /proc/meminfo | grep MemFree',shell=True)

def fn(x):
    return np.int16(x*x)

def test_plain(v):
    print "Explicit looping:"
    logmem()
    r=np.zeros(v.shape,dtype=np.int16)
    for z in xrange(v.shape[0]):
        for y in xrange(v.shape[1]):
            for x in xrange(v.shape[2]):
                r[z,y,x]=fn(x)
    print type(r[0,0,0])
    logmem()
    return r

vecfn=np.vectorize(fn,otypes=[np.int16])

def test_vectorize(v):
    print "Vectorize:"
    logmem()
    r=vecfn(v)
    print type(r[0,0,0])
    logmem()
    return r

logmem()    
s=(512,512,512)
v=np.ones(s,dtype=np.int16)
logmem()
test_plain(v)
test_vectorize(v)
v=None
logmem()

我使用的是在amd64 Debian Squeeze系统上当前的Python/numpy版本（Python 2.6.6，numpy 1.4.1）。

Answer 1

这是一个关于向量化的基本问题，所有的中间值也是向量。虽然这样做可以让程序运行得更快，但在内存使用上可能会很低效，并且会频繁地影响你的CPU缓存。为了克服这个问题，你需要使用一种方法，让循环的运行速度和编译后的代码一样快，而不是像Python那样慢。实现这个的最好方法是使用Cython，或者用f2py封装的Fortran代码，或者使用numexpr。你可以在这里找到这些方法的比较，虽然它主要关注的是速度而不是内存使用。

Answer 2

你可以查看vectorize()的源代码。它会把数组的数据类型转换成对象类型，然后调用np.frompyfunc()来根据你的Python函数创建一个ufunc（通用函数）。这个ufunc会返回一个对象数组，最后vectorize()会把这个对象数组转换成int16数组。

当数组的数据类型是对象时，会占用很多内存。

用Python函数进行逐个元素的计算速度比较慢，即使通过frompyfunc()转换成了ufunc也一样。