Оптимизация доступа к массивам numpy для numba

Недавно я наткнулся на numba и подумал о замене некоторых самодельных расширений C более элегантным кодом Python, созданным автоматически. К сожалению, я не был счастлив, когда попробовал первый быстрый тест. Кажется, что numba здесь не намного лучше, чем обычный python, хотя я ожидал почти C-подобной производительности:

from numba import jit, autojit, uint, double
import numpy as np
import imp
import logging
logging.getLogger('numba.codegen.debug').setLevel(logging.INFO)

def sum_accum(accmap, a):
    res = np.zeros(np.max(accmap) + 1, dtype=a.dtype)
    for i in xrange(len(accmap)):
        res[accmap[i]] += a[i]
    return res

autonumba_sum_accum = autojit(sum_accum)
numba_sum_accum = jit(double[:](int_[:], double[:]), 
                      locals=dict(i=uint))(sum_accum)

accmap = np.repeat(np.arange(1000), 2)
np.random.shuffle(accmap)
accmap = np.repeat(accmap, 10)
a = np.random.randn(accmap.size)

ref = sum_accum(accmap, a)
assert np.all(ref == numba_sum_accum(accmap, a))
assert np.all(ref == autonumba_sum_accum(accmap, a))

%timeit sum_accum(accmap, a)
%timeit autonumba_sum_accum(accmap, a)
%timeit numba_sum_accum(accmap, a)

accumarray = imp.load_source('accumarray', '/path/to/accumarray.py')
assert np.all(ref == accumarray.accum(accmap, a))

%timeit accumarray.accum(accmap, a)

Это дает на моей машине:

10 loops, best of 3: 52 ms per loop
10 loops, best of 3: 42.2 ms per loop
10 loops, best of 3: 43.5 ms per loop
1000 loops, best of 3: 321 us per loop

Я использую последнюю версию numba от pypi, 0.11.0. Любые предложения, как исправить код, чтобы он работал достаточно быстро с numba?


python numpy jit numba accumarray
person Michael    schedule 19.12.2013    source источник

Ответы (1)


arrow_upward
6
arrow_downward

Я разобрался сам. numba не смог определить тип результата np.max(accmap), даже если тип accmap был установлен на int. Это как-то все тормозило, но это легко исправить:

@autojit(locals=dict(reslen=uint))
def sum_accum(accmap, a):
    reslen = np.max(accmap) + 1
    res = np.zeros(reslen, dtype=a.dtype)
    for i in range(len(accmap)):
        res[accmap[i]] += a[i]
    return res

Результат довольно впечатляющий, примерно 2/3 версии C:

10000 loops, best of 3: 192 us per loop
person Michael    schedule 19.12.2013
comment
numba наиболее эффективен для чистого кода Python и массивов numpy. Никаких специальных функций он не сможет оптимально ускорить. если бы вы написали свою собственную максимальную функцию на чистом питоне и запустили ее автоматически, это, вероятно, было бы быстрее. - person M4rtini; 19.12.2013
comment
Я заглянул в код numba. Для многих функций numpy он правильно определяет тип выходных данных, просматривая входные данные и зная, как ведет себя функция. Это называется выводом типа. Это отлично работает, например. с np.sum и np.prod, но min и max не указаны в коде. Я подал отчет об ошибке для этого. - person Michael; 29.12.2013