(超级详细）jit的介绍和用法（python加速）

前言

jit是numba库有用功能里面最核心最屌的功能。

from numba import jit

普通python我不知道有没有安装numba这个库，反正anaconda是默认安装的。大家发现，numba和numpy名字有点像，没错，因为numba喜欢有numpy运算的代码，什么意思，看完本文你就知道了。

我们知道，python是解释性语言，数据类型可以是动态地，带来了很多方便，但是速度也大大降低。而编译性语言，例如c/c++很快，所以jit就是用来编译python的，编译好处就是，可以对代码进行优化，从而加速。

jit是一个修饰符decorator，作用对象是函数。即对函数进行编译优化，产生一个高效代码。

编译模式

下面介绍jit的两种编译模式。

Lazy compilation

The recommended way to use the @jit decorator is to let Numba decide when and how to optimize:

这是推荐的模式。

@jit
def f(x, y):
    # A somewhat trivial example
    return x + y

In this mode, compilation will be deferred until the first function execution. Numba will infer the argument types at call time, and generate optimized code based on this information. Numba will also be able to compile separate specializations depending on the input types. For example, calling the f() function above with integer or complex numbers will generate different code paths:

f(1, 2)#整形
#3
f(1j, 2)#复数类型
#(2+1j)

解释：jit会先让你的函数运行一次，摸清楚了传入的变量类型之后，针对这种变量类型进行优化。

Eager compilation

You can also tell Numba the function signature you are expecting. The function f() would now look like:

from numba import jit, int32

@jit(int32(int32, int32))
def f(x, y):
    # A somewhat trivial example
    return x + y

Eager急切的，渴望的。

int32(int32, int32) is the function’s signature. In this case, the corresponding specialization will be compiled by the @jit decorator, and no other specialization will be allowed. This is useful if you want fine-grained control over types chosen by the compiler (for example, to use single-precision floats).

If you omit the return type, e.g. by writing (int32, int32) instead of int32(int32, int32), Numba will try to infer it for you. Function signatures can also be strings;

解释：第一种模式是让jit自己推断数据变量的类型，而这里是你自己指定，所以优化速度明显。

注意的是，像上面，指定了数据类型都是int32，如果你拿一些其他的数据类型来，将被强制转换数据类型，因此引来精度损失或者报错。

编译选项

下面介绍如何继续提供更加精细化的控制。下面又有两个编译模式模式，nopython和object，有人会问，前面章节也是编译模式，这里又是，什么情况。这没啥的，你可以认为是对一个东西（编译）的不同角度的分类。比如人可以分成高人矮人，但也可以胖子瘦子啊，甚至男人女儿对吧，角度不一样而已。

nopython

这个模式，是被推荐的模式。

@jit(nopython=True)
def f(x, y):
    return x + y

The behaviour of the nopython compilation mode is to essentially compile the decorated function so that it will run entirely without the involvement of the Python interpreter. This is the recommended and best-practice way to use the Numba jit decorator as it leads to the best performance.

说白了就是这段代码的运行将脱离python解释器，变成机器码来运行，所以速度超快。

这个这么好，那么下面这个object模式还有什么意义呢？你错了，这个这么好是有前提的，需要你的函数代码是循环比较多，然后进行一些数学运算，这种代码在这个模式可以超级快。如果你的函数代码不是这种数学计算类的循环，比如下面这样：

def foo():
	A#非数学计算类。
	for i in range(1000):
		B#数学计算类。
	C#非数学计算类。

这个模式将自动识别那个循环，然后优化，脱离python解释器，运行。而对于A,C这两个东西无法优化，需要切换回到python解释器，极其浪费时间，效果差。切换很费时间，这种情况，最好不要用nopython的模式，而使用下面地这种普通模式。

object

普通模式，就是在python解释器里运行的模式。没有写nopython=True那么就默认是这个。

补充

问题1

从上面的描述来看，似乎已经透露了：循环，数学类的运算将大大优化，有一些代码不能优化。下面给个例子，什么东西Numba喜欢，什么类型的代码numba不喜欢。

喜欢的代码如下：

if your code is numerically orientated (does a lot of math), uses NumPy a lot and/or has a lot of loops.

例如下面的：

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))

不喜欢的就是那些numba看不懂的对象，例如你自己定义的对象object，或者一些库里面的对象，例如dataframe对象。也就是说numba喜欢搞数学有关的，以及那些基本数据类型（整形等）的运行，而对象这种东西太高级了。

例如下面这个，jit加速没有用。

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))

可以看到，这里没有使用nopython，这是对的，因为如果知道没有什么用，就不要用这个模式，否则可能加速效果适得其反。

补充一条珍贵的经验（秘密），如果随意使用nopython进行超级加速，有的时候会报错，因为你的代码让numba理解不了，很多其不知道的对象在里面。解决办法就是不要使用这个超级加速，而是用普通加速，就像上面一样。或者是不用加速（不用jit)。

问题2

使用jit之后，如何看一下是不是加速成功了？这个你自己测试一下即可。在程序运行前后都打印一下时间即可。不过要注意一个东西，就是jit第一次运行函数不会加速，因为这一次在摸清函数内部的套路。这个在前面章节我提过，怕你忘了。

所以，一个例子如下：

from numba import jit
import numpy as np
import time

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

@jit(nopython=True)
def go_fast(a): # Function is compiled and runs in machine code
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

# DO NOT REPORT THIS... COMPILATION TIME IS INCLUDED IN THE EXECUTION TIME!
start = time.time()
go_fast(x)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE
start = time.time()
go_fast(x)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))

Elapsed (with compilation) = 0.33030009269714355
Elapsed (after compilation) = 6.67572021484375e-06

可以发现，大大加速了。

一个坑

本人在使用时发现一个特别有意思的事，可能别人都不知道，算是一个告诉你的小秘密吧。那就是使用下面模式的时候

@jit(nopython=True)
def foo():

我们知道，这个编译优化，然后运行的时候会脱离python解释器。这就导致了！！！！！！！！！！！！

这一部分无法调试，你在函数内部打断点，没有用！！！！！！！！！

所以，在你调试的时候，可以先把加速给关了，不要使用jit。

参考

1.http://numba.pydata.org/numba-doc/0.17.0/user/jit.html
2.https://numba.pydata.org/numba-doc/latest/user/5minguide.html