python numba

# Python Numba让代码飞起来但别指望它是万能的最近在优化一个数值计算密集的Python项目又一次用到了Numba。这东西用好了是真能救命但用不好也真能让你头疼。今天就来聊聊这个让人又爱又恨的工具。它到底是什么Numba本质上是一个即时编译器。这么说可能有点抽象换个说法它能在运行时把Python函数翻译成机器码。这和我们平时写的Python代码运行方式完全不同——普通的Python代码是通过解释器一行行解释执行的而经过Numba编译的函数会变成直接在CPU上运行的本地代码。想象一下你有个熟练的翻译平时你说话写Python代码他就在旁边翻译给听众CPU听。现在Numba来了它直接把你要说的话提前翻译好写成稿子下次再说同样的话时直接念稿子就行。这个“写稿子”的过程就是编译念稿子的速度自然比现场翻译快得多。Numba特别擅长处理数值计算尤其是那些涉及循环、数组操作和数学运算的场景。它背后依赖的是LLVM编译器框架这个框架也被用在Clang、Rust等语言的编译器中算是编译器领域的“老将”了。它能做什么最直接的作用就是加速。但不是什么代码都能加速Numba有自己的“舒适区”。如果你有一段代码里面充满了for循环循环里又是各种数值计算那Numba很可能让这段代码快上几十倍甚至上百倍。比如计算两个大向量的点积或者对一个大矩阵进行逐元素运算这些场景都是Numba的拿手好戏。但如果你在循环里频繁调用Python对象的方法或者大量操作字符串Numba可能就帮不上什么忙了甚至可能让代码更慢。它最擅长的是那些“计算密集”而不是“数据密集”或“逻辑密集”的任务。Numba还支持GPU加速不过这个功能用起来门槛稍高一些。如果你的计算任务可以很好地映射到GPU的并行架构上比如大规模的矩阵运算那么用Numba的GPU功能可能会获得惊人的加速比。怎么用用Numba最简单的方式就是加装饰器。比如你有个计算圆周率的函数用了蒙特卡洛方法importrandomfromnumbaimportjitjit(nopythonTrue)defestimate_pi(num_samples):inside0for_inrange(num_samples):xrandom.random()yrandom.random()ifx*xy*y1.0:inside1return4.0*inside/num_samples那个jit(nopythonTrue)就是Numba的装饰器。nopythonTrue是个重要的参数它告诉Numba“必须编译成机器码如果编译不了就直接报错”。如果没有这个参数Numba在编译失败时会回退到Python模式那样就没什么加速效果了。第一次调用这个函数时会有点慢因为Numba需要时间编译。但第二次、第三次调用就会快很多。对于需要反复调用的函数这个编译时间完全可以接受。Numba还支持对NumPy数组的操作进行专门优化。如果你在函数里大量使用NumPyNumba通常能给出不错的表现。不过要注意Numba对NumPy的支持是有选择的不是所有NumPy函数都能被加速。一些实践中的经验用了这么多年Numba有些经验值得分享一下。首先不是所有代码都适合用Numba。在决定用Numba之前先用profiler找找代码的瓶颈在哪里。如果瓶颈不在数值计算上用了Numba可能也白用。其次尽量使用Numba兼容的数据类型和函数。比如用numba.typed.List代替Python原生的list用Numba提供的数学函数代替Python的math模块。Numba的文档里有个“支持的功能”列表用之前最好看看。还有注意编译开销。如果某个函数只被调用一两次而且函数本身执行时间很短那用Numba可能得不偿失——编译时间可能比执行时间还长。这种情况下要么不用Numba要么考虑预编译。调试Numba编译的函数比较麻烦。因为编译后的代码和原始Python代码已经不太一样了出错时的堆栈跟踪可能不太直观。一个实用的方法是先不用nopythonTrue让代码能运行起来确保逻辑正确后再加上这个参数进行优化。和其他技术对比说到Python的加速方案Numba有几个“兄弟”和“表亲”。最直接的对比是Cython。Cython需要你写一种类似Python但又不是Python的语言然后编译成C扩展。Cython的学习曲线比Numba陡但一旦掌握能获得更好的性能特别是对复杂项目的控制力更强。Numba的优势在于“不改动或少改动现有代码”Cython的优势在于“极致性能”。然后是PyPy。PyPy是另一个Python解释器它也有JIT编译功能。PyPy的优势是对纯Python代码友好不需要任何修改就能获得加速。但PyPy对科学计算生态的支持不如CPython很多基于C扩展的库如NumPy、Pandas在PyPy上可能有问题。还有用C/C写扩展的传统方法。这是最“重”的方法性能通常也是最好的但开发成本最高。Numba可以看作是在“方便性”和“性能”之间取了个不错的平衡点。最近一两年还有像Taichi这样的新选择。Taichi专注于图形和高性能计算语法更友好在某些领域表现比Numba更好。但生态成熟度还比不上Numba。选择哪个工具取决于你的具体需求。如果只是偶尔需要加速几段数值计算代码Numba是个不错的选择。如果需要大规模重构整个项目可能Cython更合适。如果项目对性能要求极高可能最终还是得回到C/C扩展的老路上。Numba不是银弹但它确实是个好工具。在合适的场景下它能让你用Python写出接近C速度的代码而付出的学习成本相对较低。这种“性价比”是它最大的吸引力所在。不过说到底任何工具都是这样了解它的能力边界在边界内使用才能发挥最大价值。盲目地给所有函数都加上jit装饰器往往得不到想要的结果反而可能引入新的问题。好的优化从来都是从理解问题开始的而不是从应用工具开始的。