关于使用numba对python进行代码加速

关于使用numba对python进行代码加速

前言

在处理大量数据时,pandas与numpy对单一元素的处理明显速度过慢。近期在做一篇论文的复现时,面对十万的数据集,需要做五折交叉验证,我的机器在jupyter中运行了一晚上都没有完成一半的处理,让人十分苦恼。

于是我将目光转移到各类python代码的加速方法上,常见的有Pypy,cython,numba。经过一定的了解与比较之后,我认为numba的上手难度最低,使用方法最为简单,于是将numba库引入我的代码。

现在对约50w条数据的处理用时只需8min多。

正文

虽然numba的各类文档中都提到了其使用方法很简单,调用装饰器即可,但实际使用过程中,将一篇已经完成的,使用了不少第三方库的python代码转换成经numba加速的形式,仍然会有许多困难。

函数参数的选择

  • numba是完全支持numpy的内容的,np.array可作为参数直接传入。
  • 尽量使用数字类型、numpy中的类实例作为参数,因为numba大概率不认识原生python中的数据类型(出于运算速度的考虑)。
  • 同上一点,list是不受支持的!(起初的原因是出于反射机制对性能的降低,但官方似乎已经在处理这个问题,相信在未来的numba版本中会得到改善)。
  • pandas的df类型无法作为参数传入,但由于df中的series是等同于np.array的,可以提取series转化成np.array来传参。

使用范围的界定

出于参数选择上的种种限制,用numba来重写的函数主要应该处理数值工作,如果遇到符合类型,如基于Dataframe的数据处理,需要手动将运算函数拆分出来。同时,如果要将numba运用到类成员函数上,需要参考官方文档对类进行声明,否则无法识别。

装饰器参数

@jit常用的参数有nopython,nogil,parallel,均为布尔类型。

  • nopython为True时,等同于@njit,numba会检查函数是否为非python的函数,即是否可以用c来完全重写。
  • nogil,去除GIL锁。诚然,python的低效率原因中,GIL的存在占了很大分量,numba可以让函数脱离GIL,实现真正的多线程。
  • parallel,自动并行化,此选项对矩阵运算提速有明显效果,虽然离numpy的ufunc方法仍有差距,但已经有相当的改善,同时,numba中采取vectorize装饰器写出的运算函数也可以让矩阵运算提速,但仍不及ufunc。
欢迎投喂,但你的支持就是对我最佳的回馈。