Python代码在CPU下加速：Numba入门

• Python代码加速

  主要考虑代码优化加速，而非代码逻辑优化。

  Python代码直接运行GPU是不行的，需要一定的改变，Numba是一个接口，不过本文主要针对CPU下的Python代码加速。

• Python解释器工作原理

  Python文件执行过程

  1. .py文件通过解释器转化为虚拟机可以执行的字节码(.pyc)；

  2. 字节码在虚拟机上执行，得到结果；

  字节码是一种只能运行在虚拟机上的文件，默认后缀.pyc，Python生成.pyc之后一般放在内存中继续使用，并不是每次都将.pyc文件保存到磁盘上。

  虚拟机是基于硬件和操作系统的。

  .pyc字节码通过Python虚拟机与硬件交互，就是说虚拟机的存在导致程序与硬件之间增加了中间层。效率自然被拖后。

  

• JIT（Just-In-Time）

  JIT技术中，JIT编译器将Python源代码.py直接编译成机器可以执行的机器语言（机器码），就可以直接在CPU等硬件上运行。

  这样，JIT就跳过了原来的虚拟机，执行速度几乎与用C语言编程速度无差别。

• Numba库

  Numba是Anaconda公司开发的针对Python的开源JIT编译器，用于提供Python版CPU和GPU编程，速度比原生Python快数十倍。

• 一、安装Numba

  pip install numba
  # 或者
  conda install numba
  

• 二、使用方法：装饰器

  from numba import jit
  import numpy as np
  
  SIZE = 2000
  x = np.random.random((SIZE, SIZE))
  """
  给定n*n矩阵，对矩阵每个元素计算tanh值，然后求和。
  因为要循环矩阵中的每个元素，计算复杂度为 n*n。
  """
  @jit # numba的使用方法
  def jit_tan_sum(a):   # 函数在被调用时编译成机器语言
      tan_sum = 0
      for i in range(SIZE):   # Numba 支持循环
          for j in range(SIZE):
              tan_sum += np.tanh(a[i, j])   # Numba 支持绝大多数NumPy函数
      return tan_sum
  
  print(jit_tan_sum(x))
  

  如代码所示，只是在函数上加一个装饰器@jit，就可以将运行速度提升20多倍。

  @jit装饰器的本质，是将函数编译为机器码，省掉虚拟机环节，提升速度。

• 三、Numba的使用场景总结

  numba目前只支持Python原生函数和部分Numpy函数，其他场景下无效。

  from numba import jit
  import pandas as pd
  
  x = {'a': [1, 2, 3], 'b': [20, 30, 40]}
  
  @jit
  def use_pandas(a): # Function will not benefit from Numba jit
      df = pd.DataFrame.from_dict(a) # Numba doesn't know about pd.DataFrame
      df += 1                        # Numba doesn't understand what this is
      return df.cov()                # or this!
  
  print(use_pandas(x))
  

  上述代码中使用了Pandas，而Pandas并不是原生代码，而是更高层次的封装，Numba不能理解pandas内部在做什么，所以无法对其加速。

  而一些常用的机器学习框架，比如scikit-learn, tensorflow, pyrorch等，已经做了大量的优化，不适合再使用Numba做加速。

  可以简单总结为，Numba不支持：

  • pandas

  • scikit-learn, tensorflow, pyrorch

  • try…except 异常处理

  • with 语句

  • yield from

• 四、Numba具体使用过程

  

  Numba有两种模式：

  1. @jit：object模式：上图左侧

     Numba的@jit装饰器会尝试优化代码，如果发现不支持（比如pandas等），那么Numba会继续使用Python原来的方法去执行该函数。

  2. @jit(nopython=True)或者@njit：nopython模式：上图右侧

     强制加速，不会进入上图左侧流程，只进行右侧流程，如果编译不成功，就抛出异常。

  实际使用中，一把推荐将代码中计算密集的部分作为单独的函数提出来，并使用nopython方法优化，这样可以保证能用到Numba加速；其余部分还是使用Python原生代码。

• 五、编译开销

  编译源代码需要一定的时间：

  • C/C++等编译型语言是提前把整个程序先编译好，再执行可执行文件；
  • Numba库是懒编译（Lazy Compilation）技术；

  其中，懒编译技术（Lazy Compilation）即

  1. 在运行过程中第一次发现源代码中有@jit，才将该代码块编译；
  2. 同一个Numba函数多次调用，只需要编译一次；

  $$总时间=编译时间+运行时间$$

  from numba import jit
  import numpy as np
  import time
  
  SIZE = 2000
  x = np.random.random((SIZE, SIZE))
  
  """
  给定n*n矩阵，对矩阵每个元素计算tanh值，然后求和。
  因为要循环矩阵中的每个元素，计算复杂度为 n*n。
  """
  @jit
  def jit_tan_sum(a):   # 函数在被调用时编译成机器语言
      tan_sum = 0
      for i in range(SIZE):   # Numba 支持循环
          for j in range(SIZE):
              tan_sum += np.tanh(a[i, j])   # Numba 支持绝大多数NumPy函数
      return tan_sum
  
  # 总时间 = 编译时间 + 运行时间
  start = time.time()
  jit_tan_sum(x)
  end = time.time()
  print("Elapsed (with compilation) = %s" % (end - start))
  
  # Numba将加速的代码缓存下来
  # 总时间 = 运行时间
  start = time.time()
  jit_tan_sum(x)
  end = time.time()
  print("Elapsed (after compilation) = %s" % (end - start))
  

  上述代码中，两次调用了Numba优化函数，第一次执行时需要编译，第二次就直接使用缓存的已经编译好的代码，运行时间大大缩短。

• 六、确定输入输出类型Eager Compilation：节省编译速度

  原生Python速度慢的另一个因素是变量类型不确定，Python解释器需要进行大量的类型推断。

  Numba也要推断输入输出的类型：

  from numba import jit, int32
  
  @jit("int32(int32, int32)", nopython=True)
  def f2(x, y):
      # A somewhat trivial example
      return x + y
  

  @jit(int32(int32, int32))告知Numba你的函数在使用什么样的输入和输出：

  • 括号内是输入；

  • 括号左侧是输出；

  这样不会加快执行速度，但是会加快编译速度，可以更快将函数编译到机器码上。

• Numba原理

  

  Numba使用了LLVM和NVVM技术，此技术将Python等解释型语言直接翻译成CPU、GPU可执行的机器码。

• References

1. Python | 加一行注释，让你的程序提速10+倍！numba库十分钟上手指南