众所周知,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))
得到运行时间
finish all in 0.3125
提速大约 5 倍,大成功