Numba是Python的一个JIT库,可以很方便的通过一个装饰器来实现运行时的加速。这两天初步使用了一下,有点心得体会:
1, numba喜欢for循环,有for循环的python代码(看起来像C的风格),加速效果很好;
2, numba喜欢numpy;
不过,简单的函数不要用,比如两个ndarray做减法,a-y,就不要使用@njit;
4, 如果某函数中调用的其它函数已经使用了numba的@jit或@njit,再对这个函数加速的效果就基本看不出来。
5, 还有一种不加速的情况,函数中只是调用了一个numpy函数,没有别的什么操作,这种情况,在加速和不加速的情况下,用timeit测试的结果差不多,或者只是一点点轻微的加速。(比如函数只有一句:return np.dot(a,b))
所以,如果函数中没有list comprehension,而且有for循环,可以使用@jit来加速;如果函数中全是numpy的操作,可以使用@njit来加速。
使用numba来加速代码的运行,简单上手快,比Cython好用,不过也需要考虑Python代码的可读性,不推崇一味为了加速而改写代码。使用numba,就是在速度和可读性之间寻找一个平衡位置。
下面是我自己的一段for循环用numba来加速的测试代码:
>>> from numba import jit
>>> def t1():
... return sum([x for x in range(5000)])
...
>>> @jit
... def t2():
... x = 0
... for i in range(5000):
... x += i
... return x
...
>>> timeit('t1()', number=10000, globals=globals())
3.092034629000409
>>> timeit('t1()', number=10000, globals=globals())
3.148597596002219
>>> timeit('t2()', number=10000, globals=globals())
0.11772040499636205
>>> timeit('t2()', number=10000, globals=globals())
0.0019229380050092004
>>> timeit('t2()', number=10000, globals=globals())
0.0019260349945398048
>>> timeit('t2()', number=10000, globals=globals())
0.002125131999491714
>>>
t1和t2功能完全相同,t1执行10000次3秒左右,而经过numba的@jit加速后,t2执行10000次只需要0.002秒左右,太恐怖了......用C风格写for循环,用@jit来加速,虽然有违Python的风格,但是却获得了实实在在的加速。