python 的装饰器,原理大概是将python语言转为最底层的机器码,从而实现效率的提升,对多 for 循环嵌套的情况非常有效。

@jit()
def cal(perms):
    res = List()
    tset = List([-1 for _ in range(9)])
    for perm in perms:
        for i, c in np.ndenumerate(perm):
            tset[c] = i[0]
        res.append(tset)
    return list(res)
  
  
 if __name__ == '__main__':
    perms = np.load('存着362880个array的ndarray')
    plen = len(perms)
    timestart = time.time()
    cal(perms)
    print(time.time() - timestart)
  • 未装饰时间:12.806214094161987s

  • jit装饰下时间:3.340502977371216


以下几点需要注意:
  1. 给 jit 装饰器加签名会更快,但需要输入输出为基本类型(jit reject list/set)在不签名情况下,jit 会先执行一次代码判断参数类型,第二次循环开始加速。nb.typeof(List([0]))可用来判断数据类型。
  2. jit 装饰器内只支持使用 numpy 的基本操作和一些基本数据类型:(以下是亲测出来的一些 tips)
    • np.load, np.save 是不支持的,np.array(无法指定数据类型dtype)
    • list, set 是不支持的,但有 numba.typed.List(), numba.typed.Dict()有类似功能,但均无法被直接返回,numba.typed.List()需要转换为 list 返回(注意,无法直接转换为 numpy 数组)
    • 且在加速的情况下不支持调用非传入(全局)参数,如果在 类 里或者需要使用到外部的参数,需要设置 @jit(forceobj=True),但这样无法提供加速效果。