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

小灰灰 2022-10-04 15:43 217阅读 0赞

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

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

[https://numba.readthedocs.io/en/stable/user/jitclass.html][https_numba.readthedocs.io_en_stable_user_jitclass.html]

官方文档给出一种方法 [Compiling python classes with @jitclass][Compiling python classes with _jitclass]  
比较麻烦，需要给出类中成员，而且类中如果有 list 之类的东西就会失效  
在类的外部编写函数，类中调用类外的函数，也即本文内容。  
参考 [How do I use numba on a member function of a class?][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))

得到运行时间

finish all in 0.3125

提速大约 5 倍，大成功

[https_numba.readthedocs.io_en_stable_user_jitclass.html]: https://numba.readthedocs.io/en/stable/user/jitclass.html
[Compiling python classes with _jitclass]: https://numba.pydata.org/numba-doc/0.30.1/user/jitclass.html
[How do I use numba on a member function of a class]: https://stackoverrun.com/cn/q/11493567