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

Question

以下代码最能说明我的问题:

输出到控制台(NB甚至第一次测试需要大约8分钟)显示512x512x512x16位阵列分配消耗不超过预期(每个256MByte),并且查看"顶部"过程通常仍然是600MByte如预期.

但是,在调用函数的矢量化版本时,该过程会扩展到巨大的大小(超过7GByte!).即便是我能想到的最明显的解释 - 向量化将输入和输出转换为内部的float64 - 即使向量化函数返回一个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 2.6.6,numpy 1.4.1)中最新的Python/numpy版本.

Answer 1

你可以阅读vectorize()的源代码。它将数组的 dtype 转换为对象，并调用 np.frompyfunc() 从 python 函数创建 ufunc，ufunc 返回对象数组，最后 vectorize() 将对象数组转换为 int16 数组。

当数组的数据类型为object时，会占用大量内存。

使用python函数进行逐元素计算很慢，即使是通过frompyfunc()转换为ufunc。