Python:numba 的基本应用
我们都知道Python比较慢,但很多时候我们都不知道为什么。虽然我用Python也有那么两年左右了,但也只能模模糊糊地感受到这么两点:
*Python太动态了
*如果能事先编译一下Python,让它静态一点,速度应该就会上来
于是我们就有了cython。然而cython毕竟不是原生的Python代码,使用起来还是有诸多不便的。为此,numba就成了一个功能强大又容易上手的替代选择。下面我们就先来看一下它的基本用法,在最后我们则会用卷积神经网络(CNN)的卷积和池化操作来直观感受一下其威力
(注意:以下代码的运行环境均为JupyterNotebook、Python3.6.1)
使用jit加速Python低效的for语句
jit的全称是Just-in-time,在numba里面则特指Just-in-timecompilation(即时编译)。它背后的原理我们就不细说了,总之我们看到“编译”两个字大概就能感受到它是干什么的对吧(喂
那么来看一个简单的栗子——给数组中的每个数加上一个常数c:
importnumbaasnb
importnumpyasnp
#普通的for
defadd1(x,c):
rs=[0.]*len(x)
fori,xxinenumerate(x):
rs[i]=xx+c
returnrs
#listcomprehension
defadd2(x,c):
return[xx+cforxxinx]
#使用jit加速后的for
@nb.jit(nopython=True)
defadd_with_jit(x,c):
rs=[0.]*len(x)
fori,xxinenumerate(x):
rs[i]=xx+c
returnrs
y=np.random.random(10**5).astype(np.float32)
x=y.tolist()
assertnp.allclose(add1(x,1),add2(x,1),add_with_jit(x,1))
%timeitadd1(x,1)
%timeitadd2(x,1)
%timeitadd_with_jit(x,1)
print(np.allclose(wrong_add(x,1),1))
以下是程序运行结果:
9.92ms±188μsperloop(mean±std.dev.of7runs,100loopseach)
5.77ms±347μsperloop(mean±std.dev.of7runs,100loopseach)
3.48ms±171μsperloop(mean±std.dev.of7runs,100loopseach)
需要注意的是:
*numba不支持listcomprehension,详情可参见这里::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)
使用vectorize实现numpy的Ufunc功能
虽然jit确实能让我们代码加速不少,但比之numpy的Ufunc还是要差很多:
assertnp.allclose(y+1,add_with_jit(x,1))
%timeitadd_with_jit(x,1)
%timeity+1
结果将会是:
3.76ms±292μsperloop(mean±std.dev.of7runs,100loopseach)
19.8μs±426nsperloop(mean±std.dev.of7runs,10000loopseach)
可以看到几乎有200倍的差距,这当然是无法忍受的。为此,我们可以用vectorize来定义出类似于Ufunc的函数:
@nb.vectorize(nopython=True)
defadd_with_vec(yy,c):
returnyy+c
assertnp.allclose(y+1,add_with_vec(y,1),add_with_vec(y,1.))
%timeitadd_with_vec(y,1)
%timeitadd_with_vec(y,1.)
%timeity+1
%timeity+1.
上述程序的运行结果将会是:
72.5μs±3.46μsperloop(mean±std.dev.of7runs,10000loopseach)
64.2μs±1.9μsperloop(mean±std.dev.of7runs,10000loopseach)
24.6μs±1.81μsperloop(mean±std.dev.of7runs,10000loopseach)
25.3μs±1.61μsperloop(mean±std.dev.of7runs,10000loopseach)
虽然还是慢了2倍左右,但已经好很多了
然后有几点需要注意的地方:
*vectorize下的函数所接受的参数都是一个个的数而非整个数组。所以上述add_with_vec的参数yy其实是输入数组y中的元素,而不是y本身。更详细的说明可以参见官方文档:#/numba-doc/0.23.0/user/vectorize.html)
*可以看到当常数c是整数和是浮点数时、速度是不同的。个人猜测这是因为若常数c为整数,那么实际运算时需要将它转化为浮点数,从而导致速度变慢
*上述代码中我们没有显式地定义函数的参数类型和返回类型,但我们可以预先定义。比如说,如果我确定常数c就是整数的话,我就可以这样写:
@nb.vectorize("float32(float32,int32)",nopython=True)
defadd_with_vec(yy,c):
returnyy+c
而如果我确定常数c就是浮点数的话,我就可以这样写:
@nb.vectorize("float32(float32,float32)",nopython=True)
defadd_with_vec(yy,c):
returnyy+c
而如果我确定常数c不是整数就是浮点数的话(这个人好啰嗦!),我就可以这样写:
@nb.vectorize([
"float32(float32,int32)",
"float32(float32,float32)"
],nopython=True)
defadd_with_vec(yy,c):
returnyy+c
注意,float32和float64、int32和int64是不同的,需要小心
此外,vectorize最炫酷的地方在于,它可以“并行”:
@nb.vectorize("float32(float32,float32)",target="parallel",nopython=True)
defadd_with_vec(y,c):
returny+c
assertnp.allclose(y+1,add_with_vec(y,1.))
%timeitadd_with_vec(y,1.)
%timeity+1
虽说在普通的Python3.6.1下、运行结果将如下:
73.5μs±4.22μsperloop(mean±std.dev.of7runs,10000loopseach)
21.2μs±734nsperloop(mean±std.dev.of7runs,10000loopseach)
似乎还变慢了;不过如果使用IntelDistributionforPython的话,会发现parallel版本甚至会比numpy原生的版本要稍快一些
那么是否有用parallel总会更好的栗子呢?当然是有的:
#将数组所有元素限制在某个区间[a,b]内
#小于a则置为a,大于b则置为b
#经典应用:ReLU
@nb.vectorize("float32(float32,float32,float32)",target="parallel",nopython=True)
defclip_with_parallel(y,a,b):
ify<a:
returna
ify>b:
returnb
returny
@nb.vectorize("float32(float32,float32,float32)",nopython=True)
defclip(y,a,b):
ify<a:
returna
ify>b:
returnb
returny
assertnp.allclose(np.clip(y,0.1,0.9),clip(y,0.1,0.9),clip_with_parallel(y,0.1,0.9))
%timeitclip_with_parallel(y,0.1,0.9)
%timeitclip(y,0.1,0.9)
%timeitnp.clip(y,0.1,0.9)
上述程序的运行结果将会是:
95.2μs±5.6μsperloop(mean±std.dev.of7runs,10000loopseach)
104μs±4.52μsperloop(mean±std.dev.of7runs,10000loopseach)
377μs±62μsperloop(mean±std.dev.of7runs,1000loopseach)
这个栗子中的性能提升就是实打实的了。总之,使用parallel时不能一概而论,还是要做些实验
需要指出的是,vectorize中的参数target一共有三种取值:cpu(默认)、parallel和cuda。关于选择哪个取值,官方文档上有很好的说明:
Ageneralguidelineistochoosedifferenttargetsfordifferentdatasizesandalgorithms.The“cpu”targetworkswellforsmalldatasizes(approx.lessthan1KB)andlowcomputeintensityalgorithms.Ithastheleastamountofoverhead.The“parallel”targetworkswellformediumdatasizes(approx.lessthan1MB).Threadingaddsasmalldelay.The“cuda”targetworkswellforbigdatasizes(approx.greaterthan1MB)andhighcomputeintensityalgorithms.TransferingmemorytoandfromtheGPUaddssignificantoverhead.
使用jit(nogil=True)实现高效并发(多线程)
我们知道,Python中由于有GIL的存在,所以多线程用起来非常不舒服。不过numba的jit里面有一项参数叫nogil,想来聪明的观众老爷们都猜到了它是干什么的了……
下面就来看一个栗子:
importmath
fromconcurrent.futuresimportThreadPoolExecutor
#计算类似于Sigmoid的函数
defnp_func(a,b):
return1/(a+np.exp(-b))
#参数中的result代表的即是我们想要的结果,后同
#第一个kernel,nogil参数设为了False
@nb.jit(nopython=True,nogil=False)
defkernel1(result,a,b):
foriinrange(len(result)):
result[i]=1/(a[i]+math.exp(-b[i]))
#第二个kernel,nogil参数设为了True
@nb.jit(nopython=True,nogil=True)
defkernel2(result,a,b):
foriinrange(len(result)):
result[i]=1/(a[i]+math.exp(-b[i]))
defmake_single_task(kernel):
deffunc(length,*args):
result=np.empty(length,dtype=np.float32)
kernel(result,*args)
returnresult
returnfunc
defmake_multi_task(kernel,n_thread):
deffunc(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]foriinrange(n_thread)]forarginargs]
#利用ThreadPoolExecutor进行并发
withThreadPoolExecutor(max_workers=n_thread)ase:
for_ine.map(kernel,*chunks):
pass
returnresult
returnfunc
length=10**6
a=np.random.rand(length).astype(np.float32)
b=np.random.rand(length).astype(np.float32)
b_func1=make_single_task(kernel1)
b_func2=make_multi_task(kernel1,4)
b_func3=make_single_task(kernel2)
b_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)
assertnp.allclose(rs_np,rs_nb1,rs_nb2,rs_nb3,rs_nb4)
%timeitnp_func(a,b)
%timeitnb_func1(length,a,b)
%timeitnb_func2(length,a,b)
%timeitnb_func3(length,a,b)
%timeitnb_func4(length,a,b)
这个栗子有点长,不过我们只需要知道如下两点即可:
*make_single_task和make_multi_task分别生成单线程函数和多线程函数
*生成的函数会调用相应的kernel来完成计算
上述程序的运行结果将会是:
14.9ms±538μsperloop(mean±std.dev.of7runs,100loopseach)
8.32ms±259μsperloop(mean±std.dev.of7runs,100loopseach)
10.2ms±368μsperloop(mean±std.dev.of7runs,100loopseach)
8.25ms±279μsperloop(mean±std.dev.of7runs,100loopseach)
4.68ms±114μsperloop(mean±std.dev.of7runs,100loopseach)
一般来说,数据量越大、并发的效果越明显。反之,数据量小的时候,并发很有可能会降低性能
umba的应用实例——卷积与池化
如果只想看效果的话倒没什么关系,不过如果想知道我具体在干什么的话,可以参见这篇文章:#/p/26657869
首先是卷积操作:
importnumbaasnb
importnumpyasnp
#普通的卷积
defconv_kernel(x,w,rs,n,n_channels,height,width,n_filters,filter_height,filter_width,out_h,out_w):
foriinrange(n):
forjinrange(out_h):
forpinrange(out_w):
window=x[i,...,j:j+filter_height,p:p+filter_width]
forqinrange(n_filters):
rs[i,q,j,p]+=np.sum(w[q]*window)
returnrs
#简单地加了个jit后的卷积
@nb.jit(nopython=True)
defjit_conv_kernel(x,w,rs,n,n_channels,height,width,n_filters,filter_height,filter_width,out_h,out_w):
foriinrange(n):
forjinrange(out_h):
forpinrange(out_w):
window=x[i,...,j:j+filter_height,p:p+filter_width]
forqinrange(n_filters):
rs[i,q,j,p]+=np.sum(w[q]*window)
#卷积操作的封装
defconv(x,w,kernel,args):
,n_filters=args[0],args[4]
out_h,out_w=args[-2:]
rs=np.zeros([n,n_filters,out_h,out_w],dtype=np.float32)
kernel(x,w,rs,*args)
returnrs
#64个3x28x28的图像输入(模拟mnist)
x=np.random.randn(64,3,28,28).astype(np.float32)
#16个5x5的kernel
w=np.random.randn(16,x.shape[1],5,5).astype(np.float32)
,n_channels,height,width=x.shape
_filters,_,filter_height,filter_width=w.shape
out_h=height-filter_height+1
out_w=width-filter_width+1
args=(n,n_channels,height,width,n_filters,filter_height,filter_width,out_h,out_w)
print(np.linalg.norm((conv(x,w,conv_kernel,args)-conv(x,w,jit_conv_kernel,args)).ravel()))
%timeitconv(x,w,conv_kernel,args)
%timeitconv(x,w,jit_conv_kernel,args)
上述程序的运行结果将会是:
0.00112681
3.63s±194msperloop(mean±std.dev.of7runs,1loopeach)
300ms±20.6msperloop(mean±std.dev.of7runs,1loopeach)
可以看到,仅仅是加了一个jit、速度就直接提升了十多倍
有细心的观众老爷可能已经发现,我这里没有使用np.allclose;这是因为卷积涉及到的运算太多,仅仅是将数组的dtype从float64变成float32、精度就会大幅下降,所以使用np.allclose的话会过不了assert
同时需要特别注意的是,使用jit和使用纯numpy进行编程的很大一点不同就是,不要畏惧用for;事实上一般来说,代码“长得越像C”、速度就会越快:
@nb.jit(nopython=True)
defjit_conv_kernel2(x,w,rs,n,n_channels,height,width,n_filters,filter_height,filter_width,out_h,out_w):
foriinrange(n):
forjinrange(out_h):
forpinrange(out_w):
forqinrange(n_filters):
forrinrange(n_channels):
forsinrange(filter_height):
fortinrange(filter_width):
rs[i,q,j,p]+=x[i,r,j+s,p+t]*w[q,r,s,t]
assertnp.allclose(conv(x,w,jit_conv_kernel,args),conv(x,w,jit_conv_kernel,args))
%timeitconv(x,w,jit_conv_kernel,args)
%timeitconv(x,w,jit_conv_kernel2,args)
那么程序的运行结果将会是:
281ms±12.2msperloop(mean±std.dev.of7runs,1loopeach)
66.2ms±2.32msperloop(mean±std.dev.of7runs,1loopeach)
可以看到这又快了3倍左右
接下来是池化操作(我们选用的是MaxPool):
#普通的MaxPool
defmax_pool_kernel(x,rs,*args):
,n_channels,pool_height,pool_width,out_h,out_w=args
foriinrange(n):
forjinrange(n_channels):
forpinrange(out_h):
forqinrange(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)
defjit_max_pool_kernel(x,rs,*args):
,n_channels,pool_height,pool_width,out_h,out_w=args
foriinrange(n):
forjinrange(n_channels):
forpinrange(out_h):
forqinrange(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)
defjit_max_pool_kernel2(x,rs,*args):
,n_channels,pool_height,pool_width,out_h,out_w=args
foriinrange(n):
forjinrange(n_channels):
forpinrange(out_h):
forqinrange(out_w):
_max=x[i,j,p,q]
forrinrange(pool_height):
forsinrange(pool_width):
_tmp=x[i,j,p+r,q+s]
if_tmp>_max:
_max=_tmp
rs[i,j,p,q]+=_max
#MaxPool的封装
defmax_pool(x,kernel,args):
,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)
returnrs
pool_height,pool_width=2,2
,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)
assertnp.allclose(max_pool(x,max_pool_kernel,args),max_pool(x,jit_max_pool_kernel,args))
assertnp.allclose(max_pool(x,jit_max_pool_kernel,args),max_pool(x,jit_max_pool_kernel2,args))
%timeitmax_pool(x,max_pool_kernel,args)
%timeitmax_pool(x,jit_max_pool_kernel,args)
%timeitmax_pool(x,jit_max_pool_kernel2,args)
上述程序的运行结果将会是:
586ms±38msperloop(mean±std.dev.of7runs,1loopeach)
8.25ms±526μsperloop(mean±std.dev.of7runs,100loopseach)
1.4ms±57μsperloop(mean±std.dev.of7runs,1000loopseach)
可以看到最快的比最慢的要快整整400倍有多
以上就是numba的基本应用。可能有观众想问:那有没有进阶应用?应该是有的,不过我暂时还不会,因为我也是昨天刚学……(趴