我最近偶然发现了numba,并考虑用更优雅的自动编译的python代码替换一些自制的C扩展。不幸的是,当我尝试第一个快速的基准测试时,我并不高兴。在这里,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)
在我的机器上显示:
^{pr2}$我运行的是pypi最新的numba版本,0.11.0。有什么建议,如何修复代码,使它在numba下运行得相当快?在
我自己想出来的。numba无法确定
np.max(accmap)
结果的类型,即使accmap的类型设置为int。这某种程度上减慢了一切,但修复很容易:结果相当令人印象深刻,大约是C版的2/3:
^{pr2}$相关问题 更多 >
编程相关推荐