Numba：加速python代码

Numba：加速python代码

当代码中有很多math计算，使用numpy或者有很多loops时，numba可以加速代码

最基础的是numba jit修饰器@jit

from numba import jit
import numpy as np

x = np.arange(100).reshape(10, 10)

@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
def go_fast(a): # Function is compiled to machine code when called the first time，第一次运行会被编译成机器码
    trace = 0.0
    for i in range(a.shape[0]):   # Numba likes loops
        trace += np.tanh(a[i, i]) # Numba likes NumPy functions
    return a + trace              # Numba likes NumPy broadcasting

print(go_fast(x))

nopython

@jit以两种编译模式进行，nopython编译模型实质上是编译经过修饰的函数，使得函数可以直接运行而不需要python解释器（在命令行下运行python就是启动cython解释器，cython是使用最广的python解释器）的参与。

numba工作流程

numba读取装饰函数的python字节码，并结合输入参数类型，使用LLVM编译器生成针对cpu功能量身定制的机器代码版本，然后每次调用函数会使用编译版本

numba也可以实现CUDA编程，可以使用纯python编写的内核，并让numba处理计算和数据移动

性能技巧

• loop
• 在IEEE 754合规性不太重要时，可以放宽数字上的严谨性，来获得额外的性能，实现方式为 fastmath=True
• 如果代码包含可并行化操作，numba可以编译多线程版本，实现方式为parallel=True
  • 按元素或按点数组操作
    • 一元运算符： + - ~
    • 二进制运算符： + - * / % | >> << ^ & // **
    • 比较运算符：== != < > <= >=
    • 在nopython模式下支持Numpy ufuncs
    • 用户DUFunc通过定义vectorize()
  • numpy的规约函数sum prod min max argmin argmax
  • numpy dot
  • 维度相同时，多维数组也支持以上操作，不支持混合维度数组进行numpy广播，也不支持选定维度上进行规约

支持的python功能

• 语言
  • 支持的结构
    • if ... elif ... else ...
    • while for .. in break continue
    • yield
    • assert
  • 部分受支持的结构
    • try .. except raise else finally
    • 上下文管理器
    • 列表理解
  • 不支持的构造
    • 异步
    • 类定义
    • 集合，字典，生成器理解
    • 发电机委托 yield from
• 内置类型
  • 元祖
  • 列表
  • 字典等

支持Numpy功能

• 标量类型
  • 整数
  • 布尔
  • 实数复数等
• 数组类型
  • 阵列存取
  • 属性
    • shape size ndim dtype等
  • 计算
    • all() any() argmax() argmin() mean() min()等
  • 其他方法
    • argsort()
    • astype() （仅一个参数形式）
    • copy() （无参数）
    • dot() （仅一个参数）等
  • 基本线性代数
    • 在浮点数和复数的一维和二维连续数组上支持
  • 其他功能
    • numpy.append()
    • numpy.arange()
    • numpy.ones() 等

参考文档：https://numba.readthedocs.io/en/stable/reference/numpysupported.html