Как избежать огромной дополнительной памяти потребление при использовании numpy vectorize?

Приведенный ниже код лучше всего иллюстрирует мою проблему:

Вывод на консоль (NB требуется ~ 8 минут для запуска даже первого теста) показывает, что выделение массива 512x512x512x16 бит потребляет не больше, чем ожидалось (256 МБ для каждого ), и, глядя на «верх», процесс обычно остается меньше 600 МБ, как и ожидалось.

Однако , пока вызывается векторизованная версия функции, процесс расширяется до огромного размера ( более 7 ГБ!). Даже самое очевидное объяснение, которое я могу придумать, чтобы учесть это - векторизация - это внутреннее преобразование входных и выходных данных во float64 - может составлять всего пару гигабайт, даже если векторизованная функция возвращает int16, а возвращаемый массив, безусловно, int16. Есть ли способ избежать этого? Я неправильно использую / понимаю аргумент otypes vectorize?

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()

Я использую те версии Python / numpy, которые используются в системе Amd64 Debian Squeeze (Python 2.6.6, numpy 1.4.1).

8
задан timday 16 August 2011 в 12:49
поделиться