numba中的协程

时间:2017-01-25 02:59:48

标签: python coroutine numba

我正在处理需要快速协同程序的事情,我相信numba可以加快我的代码。

这是一个愚蠢的例子:一个对其输入进行平方的函数,并将其调用的次数加到它上面。

def make_square_plus_count():
    i = 0
    def square_plus_count(x):
        nonlocal i
        i += 1
        return x**2 + i
    return square_plus_count

你甚至不能nopython=False JIT这个,大概是由nonlocal关键字引起的。

但如果您使用课程,则不需要nonlocal

def make_square_plus_count():
    @numba.jitclass({'i': numba.uint64})
    class State:
        def __init__(self):
            self.i = 0

    state = State()

    @numba.jit()
    def square_plus_count(x):
        state.i += 1
        return x**2 + state.i
    return square_plus_count

这至少有效,但如果你做nopython=True则会中断。

是否有针对nopython=True进行编译的解决方案?

1 个答案:

答案 0 :(得分:1)

如果你还要使用状态类,你也可以使用方法而不是闭包(应该是no-python编译的):

import numba

@numba.jitclass({'i': numba.uint64})
class State(object):
    def __init__(self):
        self.i = 0

    def square_plus_count(self, x):
        self.i += 1
        return x**2 + self.i

square_with_call_count = State().square_plus_count  # using the method
print([square_with_call_count(i) for i in range(10)])
# [1, 3, 7, 13, 21, 31, 43, 57, 73, 91]

然而,时间显示这实际上比纯python闭包实现慢。我希望只要您不使用nonlocal numpy-arrays或对方法(或闭包)中的数组执行操作,效率就会降低!