python - 加速用户定义的函数

Question

我有一个模拟，其中最终用户可以提供任意多个函数，然后在最内部的循环中调用这些函数。像这样的东西:

class Simulation:

    def __init__(self):
        self.rates []
        self.amount = 1

    def add(self, rate):
        self.rates.append(rate)

    def run(self, maxtime):
        for t in range(0, maxtime):
            for rate in self.rates:
                self.amount *= rate(t)

def rate(t):
    return t**2

simulation = Simulation()

simulation.add(rate)
simulation.run(100000)

作为一个Python循环，这非常慢，但我无法使用正常的方法来加速循环。

因为函数是用户定义的，所以我无法“numpyfy”最里面的调用(重写，以便最里面的工作由优化的 numpy 代码完成)。

我第一次尝试了numba，但是numba不允许将函数传递给其他函数，即使这些函数也是numba编译的。它可以使用闭包，但是因为一开始我不知道有多少个函数，所以我认为我不能使用它。关闭函数列表失败:

@numba.jit(nopython=True)
def a()
    return 1

@numba.jit(nopython=True)
def b()
    return 2

fs = [a, b]

@numba.jit(nopython=True)
def c()
    total = 0
    for f in fs:
        total += f()
    return total

c()

此操作失败并出现错误:

[...]
  File "/home/syrn/.local/lib/python3.6/site-packages/numba/types/containers.py", line 348, in is_precise
    return self.dtype.is_precise()
numba.errors.InternalError: 'NoneType' object has no attribute 'is_precise' 
[1] During: typing of intrinsic-call at <stdin> (4)

我找不到来源，但我认为 numba 的文档在某处指出这不是一个错误，但预计不会起作用。

像下面这样的东西可能可以解决从列表中调用函数的问题，但似乎是个坏主意:

def run(self, maxtime):
    len_rates = len(rates)
    f1 = rates[0]
    if len_rates >= 1:
        f2 = rates[1]
    if len_rates >= 2:
        f3 = rates[2]
    #[... repeat until some arbitrary limit]
    @numba.jit(nopython=True)
    def inner(amount):
        for t in range(0, maxtime)
            amount *= f1(t)
            if len_rates >= 1:
                amount *= f2(t)
            if len_rates >= 2:
                amount *= f3(t)
            #[... repeat until the same arbitrary limit]
        return amount

    self.amount = inner(self.amount)

我想也可以进行一些字节码黑客攻击:使用 numba 编译函数，将包含函数名称的字符串列表传递到 inner 中，执行类似 call( func_name)，然后重写字节码，使其成为 func_name(t)。

对于 cython 来说，仅编译循环和乘法可能会加速一点，但如果用户定义的函数仍然是 python，则仅调用 python 函数可能仍然会很慢(尽管我还没有对此进行分析)。我并没有真正找到关于使用 cython“动态编译”函数的太多信息，但我想我需要以某种方式向用户提供的函数添加一些类型信息，这似乎......很难。

是否有任何好的方法可以使用用户定义的函数来加速循环，而无需解析它们并从中生成代码？

Answer 1

我不认为你可以加速用户的功能 - 最终用户有责任编写高效的代码。您可以做的是提供一种以有效的方式与您的程序交互的可能性，而无需支付开销。

您可以使用 Cython，如果用户也喜欢使用 cython，那么与纯 python 解决方案相比，你们都可以实现大约 100 的加速。

作为基线，我稍微改变了你的示例:函数 rate做更多的工作。

class Simulation:

    def __init__(self, rates):
        self.rates=list(rates)
        self.amount = 1

    def run(self, maxtime):
        for t in range(0, maxtime):
            for rate in self.rates:
                self.amount += rate(t)

def rate(t):
    return t*t*t+2*t

产量:

>>> simulation=Simulation([rate])
>>> %timeit simulation.run(10**5)
43.3 ms ± 1.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

我们可以使用 cython 来加快速度，首先是你的 run功能:

%%cython
cdef class Simulation:
    cdef int amount
    cdef list rates
    def __init__(self, rates):
        self.rates=list(rates)
        self.amount = 1

    def run(self, int maxtime):
        cdef int t
        for t in range(maxtime):
            for rate in self.rates:
                self.amount *= rate(t)

这几乎给了我们因子 2:

>>> %timeit simulation.run(10**5)
23.2 ms ± 158 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

用户还可以使用 Cython 来加速计算:

%%cython
def rate(int t):
  return t*t*t+2*t

>>> %timeit simulation.run(10**5)
7.08 ms ± 145 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

使用 Cython 已经使我们的速度提高了 6，现在的瓶颈是什么？我们仍然使用 python 进行多态性/调度，这是相当昂贵的，因为为了使用它，必须创建 Python 对象(即这里的 Python 整数)。我们可以用 Cython 做得更好吗？是的，如果我们为传递给 run 的函数定义一个接口(interface)在编译时:

%%cython   
cdef class FunInterface:
   cpdef int calc(self, int t):
      pass

cdef class Simulation:
    cdef int amount
    cdef list rates

    def __init__(self, rates):
        self.rates=list(rates)
        self.amount = 1

    def run(self, int maxtime):
        cdef int t
        cdef FunInterface f
        for t in range(maxtime):
            for f in self.rates:
                self.amount *= f.calc(t)

cdef class  Rate(FunInterface):
    cpdef int calc(self, int t):
        return t*t*t+2*t

这会额外加速 7 倍:

 simulation=Simulation([Rate()])
 >>>%timeit simulation.run(10**5)
 1.03 ms ± 20.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

上面代码中最重要的部分是行:

self.amount *= f.calc(t)

不再需要 python 进行调度，而是使用与 C++ 中的虚拟函数非常相似的机制。这种 C++ 方法只有非常小的一次间接/查找开销。这也意味着函数的结果和参数都不必转换为 Python 对象。为此，Rate一定是cpdef-function，你可以看看here有关更多详细信息，请参阅 cpdef 函数的继承如何工作。

现在的瓶颈是 for f in self.rates 行因为我们仍然需要在每一步中进行大量的 python 交互。下面是一个示例，如果我们能够改进的话，可能会实现以下效果:

%%cython
.....
cdef class Simulation:
    cdef int amount
    cdef FunInterface f  #just one function, no list

    def __init__(self, fun):
        self.f=fun
        self.amount = 1

    def run(self, int maxtime):
        cdef int t
        for t in range(maxtime):
                self.amount *= self.f.calc(t)

...

 >>>  simulation=Simulation(Rate())
 >>> %timeit simulation.run(10**5)
 408 µs ± 1.41 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

另一个因素2，但您可以决定是否需要更复杂的代码，以便存储 FunInterface 的列表。 -没有 python 交互的对象确实值得。