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

Question

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

Вывод на консоль (примечание: для запуска даже первого теста требуется ~ 8 минут) показывает, что 512x512x512x16-битные выделения массива занимают не больше ожидаемого (256MByteдля каждого) и, глядя на «верх», процесс, как правило, остается менее 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, которые используются в системе Debian Squeeze amd64 (Python 2.6.6, numpy 1.4.1).

Answer 1

Основная проблема векторизации состоит в том, что все промежуточные значения также являются векторами. Хотя это удобный способ получить приличное повышение скорости, он может быть очень неэффективным с использованием памяти и будет постоянно перегружать кэш вашего процессора. Чтобы преодолеть эту проблему, вам нужно использовать подход, который имеет явные циклы, работающие на скорости компиляции, а не на скорости Python. Лучший способ сделать это - использовать cython , код Fortran, завернутый в f2py или Numberxpr . Вы можете найти сравнение этих подходов здесь , хотя здесь больше внимания уделяется скорости, чем использованию памяти.

Answer 2

Вы можете прочитать исходный код vectorize ().Он преобразует dtype массива в объект и вызывает np.frompyfunc () для создания ufunc из вашей функции python, ufunc возвращает массив объектов и, наконец, vectorize () преобразует массив объектов в массив int16.

Он будет использовать много памяти, когда массив типа d - объект.

Использование функции python для поэлементного вычисления медленное, даже преобразовывается в ufunc с помощью pypyunc ().