在类中使用numba提高python运行速度

原创

宏睿时空 2021-07-13 11:04:51 ©著作权

©著作权归作者所有：来自51CTO博客作者宏睿时空的原创作品，请联系作者获取转载授权，否则将追究法律责任

众所周知，numba 的 jit 功能可以显著提高 python 的运行速度。
然而，numba 对于程序有一些特定的要求，比如不能用 list 之类的，而且似乎也不能直接放在类（class）里。

要想在类中加速的话，有两种方法：

https://numba.readthedocs.io/en/stable/user/jitclass.html

官方文档给出一种方法 Compiling python classes with @jitclass
比较麻烦，需要给出类中成员，而且类中如果有 list 之类的东西就会失效
在类的外部编写函数，类中调用类外的函数，也即本文内容。
参考 How do I use numba on a member function of a class?
解决方法
对于以下代码（没有 jit 修饰）：

class A():
    def __init__(self):
        self.a = 1
        self.n = 10000000

    def calc(self):
        for _ in range(self.n):
            self.a += 1


if __name__ == '__main__':
    start = time.process_time()
    a = A()
    a.calc()
    end = time.process_time()
    print('finish all in %s' % str(end - start))

正常运行，得到运行时间

finish all in 1.546875

若想加速 calc() 函数，直接在类中加上 jit 修饰，会报 warning，并且 numba 退化

    @jit
    def calc(self):
        for _ in range(self.n):
            self.a += 1

改写如下形式：

@jit
def _calc(n, a):  # 在外部写一个函数
    for _ in range(n):
        a += 1
    return a


class A():
    def __init__(self):
        self.a = 1
        self.n = 10000000

    def calc(self):  # 类中调用外部的函数
        self.a = _calc(self.n, self.a)


if __name__ == '__main__':
    start = time.process_time()
    a = A()
    a.calc()
    end = time.process_time()
    print('finish all in %s' % str(end - start))

得到运行时间