使用numba加速python程序

前言

最近在写一个实时图像处理的程序，程序需要对快速相机捕获的目标图像进行实时地计算处理并输出处理结果。由于相机帧频比较快（几百fps），计算机对图像的处理速度就成了整个程序速度的瓶颈，在使用python语言编写完成后，程序的运行速度只有几十赫兹。为了提高程序的运行速度，引入了numba工具，使用jit和vectorize两种函数修饰器对程序中比较耗时的运算进行修饰后，程序的运行速度达到几百赫兹，接近相机的帧频。
由于程序中图像的分辨率不是特别高，计算的并行读较低，对计算速度的要求也不是特别高，numba的jit和vectorize这两种简单方便的修饰器正好能够满足我的需求。下面介绍一下这两种修饰器的使用方法。

准备

开始之前，需要先安装相关的工具包，主要包括numpy、llvmlite、numba。
如果你的电脑连接了互联网的话，这些工具包的安装都非常简单。如果你安装了Anaconda的conda工具的话，只需要在cmd命令窗口输入conda install numba即可，如果没有conda的话也可以通过pip工具完成安装：pip install numba。
如果你的电脑没有连接互联网的话，安装起来就会比较麻烦了，你需要先从这里下载跟你的操作系统和python版本匹配的安装包（llvmlite、numpy、numba），然后再使用pip工具依次安装：pip install path\filename.whl
安装完成后新建一个py文件，输入：

import numpy as np
from numba import jit, vectorize

运行没有问题的话说明安装成功了。

jit修饰器的使用

先来看下面这段程序：

import numpy as np
from numba import jit,vectorize
import time

 #不使用jit函数修饰器的函数
def func1_without_jit(A,C):
    for i in range(30):
        for j in range(30):
            grid=A[10*i:10*(i+1),10*j:10*(j+1)]
            C[i,j]=grid.sum()
    return C
    
 #使用jit函数修饰器   
@jit(nopython=True,nogil=True)
def func1_with_jit(A,C):
    for i in range(30):
        for j in range(30):
            grid = A[10*i:10*(i+1),10*j:10*(j+1)]
            C[i,j]=grid.sum()
    return C
    
 #使用jit函数修饰器 
@jit(nopython=True,nogil=True)
def func1_with_jit1(A,C):
    for i in range(30):
        for j in range(30):
            grid = A[10*i:10*(i+1),10*j:10*(j+1)]
            C[i,j]=0
            for m in range(10):
                for n in range(10):
                    C[i,j]+=grid[m,n]
    return C


if __name__=='__main__':
    A=np.random.rand(90000).astype(np.float32)
    A=A.reshape(300,300)
    C=np.zeros([30,30],dtype=np.float32)
    #先对三个函数预编译一次
    C=func1_without_jit(A,C)
    C=func1_with_jit(A,C)
    C = func1_with_jit1(A, C)
    
    #计算func1_without_jit运行5000次的时间
    t1=time.time()
    for i in range(5000):
        C=func1_without_jit(A,C)
    print('Time without jit is:'+str(time.time()-t1))
    
   #计算func1_with_jit运行5000次的时间 
    t1 = time.time()
    for i in range(5000):
        C = func1_with_jit(A, C)
    print('Time with jit is:' + str(time.time() - t1))
    
     #计算func1_with_jit1运行5000次的时间 
    t1 = time.time()
    for i in range(5000):
        C = func1_with_jit1(A, C)
    print('Time with jit1 is:' + str(time.time() - t1))

这段程序的作用是将一个300×300的矩阵压缩成30×30，新矩阵的每个点是原来矩阵对应区域的和。我们定义了三个函数来实现这一功能，其中第一个没有使用jit修饰器，后面两个都使用了jit修饰器，所不同的是，第三个函数写得更‘C语言化’，将np.sum函数也展开使用for循环替代。这段程序的运行结果如下所示：

D:/Python/NumbaTest/NumbaTest.py
Time without jit is:17.769345998764038
Time with jit is:0.7519514560699463
Time with jit1 is:0.5844631195068359

Process finished with exit code 0

使用jit修饰器后，函数的运行时间缩短了20~30倍。
从这段程序的运行结果可以看出：
1、jit修饰器的使用方法很简单，在定义函数前增加一行声明@jit(nopython=True,nogil=True)即可，括号里的参数设置是可选内容，具体描述可见链接，这里的nopython表示不使用python编译器，nogil表示不使用python的gil全局锁，只有在nopython模式下才能使用nogil模式。
2、通常情况下jit修饰器对for循环的加速效果很好，使用jit修饰器时，不用害怕使用多重循环，函数写的越像C语言，加速效果越好。
3、jit修饰的函数在第一次调用时会对其进行jit编译并将编译结果记录下来以供后续调用使用，因此第一次调用速度较慢，这也是为什么程序里先调用一次后再计时的原因。
4、需要注意的是jit目前只能支持一部分的numpy函数，如果调用了不支持的函数程序会报错。

vectorize修饰器

下面来看vectorize修饰器，其使用方法与jit类似。vectorize修饰器适合对矩阵之间的并行操作进行加速。
看下面的代码：

def func2_without_vec(A,B,C):
    return (A-B*C)*(B>0.5)*(C>0.5)
@vectorize(nopython=True)
def func2_with_vec(A,B,C):
    return (A-B*C)*(B>0.5)*(C>0.5)

if __name__=='__main__':
    A=np.random.rand(90000).astype(np.float32)
    A=A.reshape(300,300)
    B = np.random.rand(90000).astype(np.float32)
    B = B.reshape(300, 300)
    C = np.random.rand(90000).astype(np.float32)
    C = B.reshape(300, 300)

    func2_without_vec(A, B,C)
    func2_with_vec(A, B,C)

    t1=time.time()
    for i in range(5000):
        func2_without_vec(A, B,C)
    print('Time without vectorize is:' + str(time.time() - t1)+'s')

    t1 = time.time()
    for i in range(5000):
        func2_with_vec(A, B,C)
    print('Time with vectorize is:' + str(time.time() - t1) + 's')

程序的运行结果为：

Time without vectorize is:2.327772378921509s
Time with vectorize is:0.14960002899169922s

Process finished with exit code 0

这里需要注意的是A,B,C必须是同纬度的矩阵或者数字，函数func2_with_vec(A,B,C)对矩阵的操作必须可以拆分成对每一个元素的操作。例如：A+B可以等效为

A[i,j]+B[i,j]，for each element pairs in A and B

因此A+B可以使用vectorize加速，而像矩阵的转置、求逆、矩阵的叉乘等运算则不能使用vectorize加速，否则会报错。