Python高性能计算库—Numba
摘要： 在计算能力为王的时代，具有高性能计算的库正在被大家应用于深度学习。例如：Numpy，本文介绍了一个新的Python库——Numba， 在计算性能方面，它比Numpy表现的更好。 最近我在观看一些SciPy2017会议的视频，偶然发现关于Numba的来历--讲述了那些C++ 横行者因为对Gil Forsyth和Lorena Barba失去信心而编写的一个库。虽然本人觉得这个想法有些不妥，但我真的很喜欢他们所教授的知识。因为我发现自己正在受益于这个库，并且从Python代码中获得了令人难以置信的表现，所以我觉得应该要写一些关于Numba库的介绍性文章，也可能会在将来添加一系列小的更多类似教程的文章。
1.那么到底什么是Numba？
Numba是一个库，可以在运行时将Python代码编译为本地机器指令，而不会强制大幅度的改变普通的Python代码（稍后再做说明）。翻译/魔术是使用LLVM编译器完成的，该编译器是相当活跃的开源社区开发的。 Numba最初是由ContinuumAnalytics内部开发，此公司也开发了著名的Anaconda，但现在它是开源的。核心应用领域是 math-heavy（密集数学？重型数学？）和 array-oriented（面向数组）的功能，它们在本地Python中相当缓慢。想象一下，在 Python中编写一个模块，必须一个元素接着一个元素的循环遍历一个非常大的数组来执行一些计算，而不能使用向量操作来重写。这是很不好的主意，是吧？所以“通常”这类库函数是用C/C++或Fortran编写的，编译后，在Python中作为外部库使用。Numba 这类函数也可以写在普通的Python模块中，而且运行速度的差别正在逐渐缩小。
2.怎么才能get到Numba呢？
安装Numba的推荐方法是使用conda包管理 conda install numba 你也可以用 pip 来安装Numba，但是最新版本的发布才一天之久。但是，只要你能够使用 conda，我会推荐使用它，因为它能够为你安装例如CUDA工具包，也许你想让你的 Python代码GPU就绪（当然，这也是有可能的！）。
3.如何使用Numba呢？
@jit 使用它的要求不多。基本上，你写一个自己的“普通”的Python函数，然后给函数定义添加一个装饰（如果你不是很熟悉装饰器，读一下关于this或that）。你可以使用不同类型的装饰器，但可能是刚开始的选择之一。其他装饰器可用于例如创建 @vectorize 或编写将在CUDA GPU上执行的代码 @cuda numpy通用功能
基本的步骤。他们提供的代码示例是2d数组的求和函数，以下是代码：
In [2]:
import numba as nb
import numpy as np
# 普通的 for
def add1(x, c):
rs = [0.] * len(x)
for i, xx in enumerate(x):
rs[i] = xx + c
return rs
# list comprehension
def add2(x, c):
return [xx + c for xx in x]
# 使用 jit 加速后的 for
@nb.jit(nopython=True)
def add_with_jit(x, c):
rs = [0.] * len(x)
for i, xx in enumerate(x):
rs[i] = xx + c
return rs
y = np.random.random(10**5).astype(np.float32)
x = y.tolist()
assert np.allclose(add1(x, 1), add2(x, 1), add_with_jit(x, 1))
%timeit add1(x, 1)
%timeit add2(x, 1)
%timeit add_with_jit(x, 1)
print(np.allclose(add_with_jit(x, 1), 1))
13.8 ms ± 1.72 ms per loop (mean ± std. dev. of 7 runs, 100 loops each) 8.03 ms ± 369 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 4.88 ms ± 320 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) False
需要注意的是：
-
numba不支持 list comprehension，详情可参见这里
-
jit能够加速的不限于for，但一般而言加速for会比较常见、效果也比较显著。我在我实现的numpy版本的卷积神经网络（CNN）中用了jit后、可以把代码加速 20 倍左右。具体代码可以参见这里，不过如果不想看源代码的话，可以参见CNN.ipynb，我在其中做了一些相应的、比较简单的实验
-
jit会在某种程度上“预编译”你的代码，这意味着它会在某种程度上固定住各个变量的数据类型；所以在jit下定义数组时，如果想要使用的是float数组的话，就不能用[0] len(x)定义、而应该像上面那样在0后面加一个小数点：[0.] len(x)
In [3]:
@nb.jit
def f():
return [x * 2 for x in range(5)]
f()
Out[3]:
[0, 2, 4, 6, 8]
In [10]:
# from numba import nb
# from numpy import np
# jit decorator tells Numba to compile this function.
# The argument types will be inferred by Numba when function is called.
@nb.jit
def sum2d(arr):
M, N = arr.shape
result = 0.0
for i in range(M):
for j in range(N):
result +=arr[i,j]
return result
a = np.arange(9).reshape(3,3)
print (sum2d(a))
36.0
正如你所看到的，Numba 装饰器被添加到函数定义中，并且 voilá这个函数将运行得很 快。但是，这里带来了很有趣的注意事项：你只能使用 Numpy 和标准库里的函数来加快 Numba 速度，甚至不需要开了他们所有的特性。他们有一个相当好的文档（参考资料）， 列出了所有支持的内容。见 here 是所支持 Python 的功能和 here 是所支持的 Numpy 功 能。现在支持的功能可能还不太多，但我想告诉你，这就够了！请记住，Numba 不是要 加快你的数据库查询或如何强化图像处理功能。他们的目标是加快面向数组的计算，我 们可以使用它们库中提供的函数来解决。
4. 示例和速度比较
熟练的 Python 用户永远不会使用上述代码实现 sum 功能，而是调用 numpy.sum。相反， 我将向你介绍另外一个例子，为了更好地理解这个例子，也许刚开始是一个小的背景故 事（如果你对这个例子的背景不感兴趣，你可以直接跳过然后直接去看代码）。 从我所学习的知识来看，我会认为自己是一个水文学家，我做的很多的一件事是模拟降 雨径流过程。简单点来说：通过时间序列数据，例如雨量和空气温度，然后尝试创建模 型来判断一条河流的水流量有多少。这在外行看来是非常复杂。但，对于我们来说，很 简单。我们通常使用的模块迭代输入数组，并且对于每个时间步长，我们会更新一些模 块内部的状态（例如，模拟土壤水分，积雪或拦截水中的树木）。在每个时间段结束时， 计算水流量，这不仅取决于在同一时间步长下的雨，而且也取决于在内部模型状态（或 储存）。在这种情况下，我们就需要考虑以前时间步长的状态和输出。那么你可能会看 到这个问题：我们必须一段时间接一段时间的计算整个流程，而对于解决这种问题 Python 本来就是很慢的！这就是为什么大多数模块都是在 Fortran 或 C/C ++中实现的。 如前所述：Python 在对于这种面向数组的计算来说是慢的。但是 Numba 允许我们在 Python 中做同样的事情，而且没有太多的性能损失。我认为至少对于模型的理解和发展， 这可能会很方便。（所以我最近创建了一个名为“RRMPG ”的项目——降雨径流建模游 乐场）。 Okay，现在我们来看看我们 get 到了什么。我们将使用最简单的模块之一，由 MB Fiering 在1967年出于教育目的开发的ABC模型，并将Python代码的速度与Numba优化后Python 代码和 Fortran 实现进行比较。请注意这个模型不是我们在现实中使用的（正如名称所 示），但是我认为这可能是一个不错的想法来举例。 ABC 模块是一个三个参数模块（a，b，c，习惯性命名），它只接收下雨量为输入，只有 一个存储。土壤水分蒸发蒸腾损失总量（参数 b），另一部分通过土壤渗透到地下水储 存（参数 a），最后一个参数 c 代表地下水总量，离开地下变成河流。Python 中的代码， 使用 Numpy 数组可能会像如下所示：
In [14]:
# import numpy as np
def abc_model_py(a, b, c, rain):
# initialize array for the stream discharge of each time step
outflow = np.zeros((rain.size), dtype=np.float64)
# placeholder, in which we save the storage content of the previous and
# current timestep
state_in = 0
state_out = 0
for i in range(rain.size):
# Update the storage
state_out = (1 - c) * state_in + a * rain[i]
# Calculate the stream discharge
outflow[i] = (1 - a - b) * rain[i] + c * state_out
state_in = state_out
return outflow
In [16]:
@nb.jit
def abc_model_numba(a, b, c, rain):
outflow = np.zeros((rain.size), dtype=np.float64)
state_in = 0
state_out = 0
for i in range(rain.size):
state_out = (1 - c) * state_in + a * rain[i]
outflow[i] = (1 - a - b) * rain[i] + c * state_out
state_in = state_out
return outflow
In [21]:
rain = np.arange(9) py_time = %timeit -r 5 -n 10 -o abc_model_py(0.2, 0.6, 0.1, rain) py_time
225 µs ± 28.7 µs per loop (mean ± std. dev. of 5 runs, 10 loops each)
Out[21]:
<TimeitResult : 225 µs ± 28.7 µs per loop (mean ± std. dev. of 5 runs, 10 loops each)>
In [22]:
numba_time = %timeit -r 5 -n 10 -o abc_model_numba(0.2, 0.6, 0.1, rain) numba_time
The slowest run took 12580.76 times longer than the fastest. This could mean that an intermediate result is being cached. 2.26 ms ± 4.52 ms per loop (mean ± std. dev. of 5 runs, 10 loops each)
Out[22]:
<TimeitResult : 2.26 ms ± 4.52 ms per loop (mean ± std. dev. of 5 runs, 10 loops each)>
In [23]:
py_time.best / numba_time.best
Out[23]:
220.61904702829545
通过添加一个装饰器，我们的计算速度比纯Python代码快220.61904702829545倍, 利用了JIT技术，你调用函数实际上是经过llvm编译过的，编译的过程应当是
python 字节码 -> llvm IR -> 二进制，利用装饰器重定向了函数的调用，使用ctypes调用编译后的代码。
当时对jit技术感兴趣时，发现了这样一个有趣的repo https://github.com/cslarsen/minijit
实现了一个简单的类似numba的jit编译器。贴一下装饰器的实现思路，具体如何用ctypes定位到编译后的代码，楼主可以查看该repo的实现。
In [24]:
def jit(function):
def frontend(*args, **kw):
if not hasattr(frontend, "function"):
try:
# numba中使用llvm将字节码编辑到二进制，该repo中自己实现了
native, asm = compile_native(function, verbose=False)
native.raw = asm.raw
native.address = asm.address
# 将编译后的函数存在编译函数对象的function中，
frontend.function = native
except Exception as e:
frontend.function = function # fallback to Python
return frontend.function(*args, **kw)
return frontend
In [ ]: