numba 笔记

目录

numba为Python提速

numba能够极大的提高python在计算方面的性能。是不是所有的python代码上，都可以加上numba.jit装饰器？答案是否定的。

同时需要特别注意的是，使用jit和使用纯numpy进行编程的很大一点不同就是，不要畏惧用for；事实上一般来说，代码“长得越像 C”、速度就会越快：

使用 jit(nogil=True) 实现高效并发（多线程）

 

使用Numba

要求不多。基本上，你写一个自己的“普通”的Python函数，然后给函数定义添加一个装饰（如果你不是很熟悉装饰器，读一下关于this或that）。你可以使用不同类型的装饰器，但@jit可能是刚开始的选择之一。其他装饰器可用于例如创建numpy通用功能@vectorize或编写将在CUDA GPU上执行的代码@cuda。
 

 

numba为Python提速

注释掉@jit，foo()耗时17.66000986099243秒，foo2()耗时6.496371507644653秒；
启用@jit，foo()耗时0.04500269889831543秒，foo2()耗时0.041002511978149414秒；
foo()相差392倍，foo2()相差158倍；
@jit必须在要加速的函数上方（不能隔其他函数）才有效，可同时使用多个@jit。

注意下面的情况：

@jit								
def foo():
	x = []
	for a in range(10000000):
		x.append(a)

def foo2():	
	y = []					   
	for b in range(10000000):
		y.append(b)

if __name__ == '__main__':
	start = time()
	foo()
	end = time()
	shi = end - start
	start2 = time()
	foo2()
	end2 = time()
	shi2 = end2 - start2
	print("foo()总耗时：%s秒" % shi)
	print("foo2()总耗时：%s秒" % shi2)

当遍历100000次时：

foo()总耗时：0.11800670623779297秒
foo2()总耗时：0.011000633239746094秒

当遍历1000000次时：

foo()总耗时：0.1390078067779541秒
foo2()总耗时：0.1480085849761963秒

当遍历10000000次时：

foo()总耗时：0.5130293369293213秒
foo2()总耗时：1.4940855503082275秒

当遍历100000000次时：电脑崩溃了，内存占用太大。
综上：@jit适合数据量大时使用，数据量少时尽量不要使用。

 

 

numba能够极大的提高python在计算方面的性能。是不是所有的python代码上，都可以加上numba.jit装饰器？答案是否定的。

示例
环境

python3.6
fedora
pymysql

示例
很常见的例子，从数据库从查询一千条数据，再进行简单的格式转换。

#coding=utf-8
import time
from numba import jit

import pymysql
conn = pymysql.connect("192.168.10.125", "username", "password", "db_name")

#@jit   #位置1
def test_for():
    c =conn.cursor()
    c.execute("select * from user limit 1000")

    start = time.time()
    for a in c.fetchall():
        to_dict(a, c.description)
    end = time.time()
    print("cost_seconds:", end-start)

#@jit   #位置2
def to_dict(row, description):
    """
    记录转为字典
    """
    item = dict()
    for i, field in enumerate(description):
        field_name = field[0]
        item[field_name] = row[i]
    return item

test_for()

结果
python3           0.0033788681030273438
python3 +位置1jit 0.08921098709106445
python3 + 位置1jit+位置2jit  3.846426010131836
pypy              0.0065860748291015625

 

同时需要特别注意的是，使用jit和使用纯numpy进行编程的很大一点不同就是，不要畏惧用for；事实上一般来说，代码“长得越像 C”、速度就会越快：

 

池化操作（我们选用的是 MaxPool）：

# 普通的 MaxPool
def max_pool_kernel(x, rs, *args):
    n, n_channels, pool_height, pool_width, out_h, out_w = args
    for i in range(n):
        for j in range(n_channels):
            for p in range(out_h):
                for q in range(out_w):
                    window = x[i, j, p:p+pool_height, q:q+pool_width]
                    rs[i, j, p, q] += np.max(window)

# 简单地加了个 jit 后的 MaxPool
@nb.jit(nopython=True)
def jit_max_pool_kernel(x, rs, *args):
    n, n_channels, pool_height, pool_width, out_h, out_w = args
    for i in range(n):
        for j in range(n_channels):
            for p in range(out_h):
                for q in range(out_w):
                    window = x[i, j, p:p+pool_height, q:q+pool_width]
                    rs[i, j, p, q] += np.max(window)

# 不惧用 for 的、“更像 C”的 MaxPool
@nb.jit(nopython=True)
def jit_max_pool_kernel2(x, rs, *args):
    n, n_channels, pool_height, pool_width, out_h, out_w = args
    for i in range(n):
        for j in range(n_channels):
            for p in range(out_h):
                for q in range(out_w):
                    _max = x[i, j, p, q]
                    for r in range(pool_height):
                        for s in range(pool_width):
                            _tmp = x[i, j, p+r, q+s]
                            if _tmp > _max:
                                _max = _tmp
                    rs[i, j, p, q] += _max

# MaxPool 的封装
def max_pool(x, kernel, args):
    n, n_channels = args[:2]
    out_h, out_w = args[-2:]
    rs = np.zeros([n, n_filters, out_h, out_w], dtype=np.float32)
    kernel(x, rs, *args)
    return rs

pool_height, pool_width = 2, 2
n, n_channels, height, width = x.shape
out_h = height - pool_height + 1
out_w = width - pool_width + 1
args = (n, n_channels, pool_height, pool_width, out_h, out_w)

assert np.allclose(max_pool(x, max_pool_kernel, args), max_pool(x, jit_max_pool_kernel, args))
assert np.allclose(max_pool(x, jit_max_pool_kernel, args), max_pool(x, jit_max_pool_kernel2, args))
%timeit max_pool(x, max_pool_kernel, args)
%timeit max_pool(x, jit_max_pool_kernel, args)
%timeit max_pool(x, jit_max_pool_kernel2, args)

上述程序的运行结果将会是：

586 ms ± 38 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
8.25 ms ± 526 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1.4 ms ± 57 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

可以看到最快的比最慢的要快整整 400 倍有多

 

 

* numba不支持 list comprehension，详情可参见这里::https://github.com/numba/numba/issues/504

* jit能够加速的不限于for，但一般而言加速for会比较常见、效果也比较显著。我在我实现的numpy版本的卷积神经网络（CNN）中用了jit后、可以把代码加速 60 多倍。具体代码可以参见这里:https://github.com/carefree0910/MachineLearning/blob/master/NN/Basic/Layers.py#L9

* jit会在某种程度上“预编译”你的代码，这意味着它会在某种程度上固定住各个变量的数据类型；所以在jit下定义数组时，如果想要使用的是float数组的话，就不能用[0] * len(x)定义、而应该像上面那样在0后面加一个小数点：[0.] * len(x)

 

 

使用 jit(nogil=True) 实现高效并发（多线程）

我们知道，Python 中由于有 GIL 的存在，所以多线程用起来非常不舒服。不过 numba 的 jit 里面有一项参数叫 nogil，想来聪明的观众老爷们都猜到了它是干什么的了……

下面就来看一个栗子：

import math
from concurrent.futures import ThreadPoolExecutor

# 计算类似于 Sigmoid 的函数
def np_func(a, b):
    return 1 / (a + np.exp(-b))

# 参数中的 result 代表的即是我们想要的结果，后同
# 第一个 kernel，nogil 参数设为了 False
@nb.jit(nopython=True, nogil=False)
def kernel1(result, a, b):
    for i in range(len(result)):
        result[i] = 1 / (a[i] + math.exp(-b[i]))

# 第二个 kernel，nogil 参数设为了 True
@nb.jit(nopython=True, nogil=True)
def kernel2(result, a, b):
    for i in range(len(result)):
        result[i] = 1 / (a[i] + math.exp(-b[i]))

def make_single_task(kernel):
    def func(length, *args):
        result = np.empty(length, dtype=np.float32)
        kernel(result, *args)
        return result
    return func

def make_multi_task(kernel, n_thread):
    def func(length, *args):
        result = np.empty(length, dtype=np.float32)
        args = (result,) + args
        # 将每个线程接受的参数定义好
        chunk_size = (length + n_thread - 1) // n_thread
        chunks = [[arg[i*chunk_size:(i+1)*chunk_size] for i in range(n_thread)] for arg in args]
        # 利用 ThreadPoolExecutor 进行并发
        with ThreadPoolExecutor(max_workers=n_thread) as e:
            for _ in e.map(kernel, *chunks):
                pass
        return result
    return func

length = 10 ** 6
a = np.random.rand(length).astype(np.float32)
b = np.random.rand(length).astype(np.float32)

nb_func1 = make_single_task(kernel1)
nb_func2 = make_multi_task(kernel1, 4)
nb_func3 = make_single_task(kernel2)
nb_func4 = make_multi_task(kernel2, 4)

rs_np = np_func(a, b)
rs_nb1 = nb_func1(length, a, b)
rs_nb2 = nb_func2(length, a, b)
rs_nb3 = nb_func3(length, a, b)
rs_nb4 = nb_func4(length, a, b)
assert np.allclose(rs_np, rs_nb1, rs_nb2, rs_nb3, rs_nb4)
%timeit np_func(a, b)
%timeit nb_func1(length, a, b)
%timeit nb_func2(length, a, b)
%timeit nb_func3(length, a, b)
%timeit nb_func4(length, a, b)

这个栗子有点长，不过我们只需要知道如下两点即可：

* make_single_task和make_multi_task分别生成单线程函数和多线程函数

* 生成的函数会调用相应的kernel来完成计算

上述程序的运行结果将会是：

14.9 ms ± 538 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
8.32 ms ± 259 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
10.2 ms ± 368 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
8.25 ms ± 279 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.68 ms ± 114 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

一般来说，数据量越大、并发的效果越明显。反之，数据量小的时候，并发很有可能会降低性能