利用numba給Python代码加速 [2]

2022-01-10 08:09:09 浏览数 (1)

  • @vectorize 装饰器

Numba 的 @vectorize 装饰器可以将以标量为输入的的python函数编译为类似Numpy的 ufuncs。创建一个传统的NumPy ufunc并不是最简单的过程,它可能需要编写一些C代码。Numba让这很容易。使用@vectorize装饰器 ,Numba可以将纯Python函数编译成ufunc,该ufunc在NumPy数组上运行的速度与用C编写的传统ufunc一样快。

一个简单的例子:

代码语言:javascript复制
import numpy as np
from numba import vectorize, float64

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

a = np.arange(1,11)
b = np.arange(10,0,-1)
print(f(a,b))

以上是惰性编译,不指定参数类型。也可以采用积极编译,给出函数签名,即指定返回值和输入参数的类型。注意函数签名需写在列表中。

代码语言:javascript复制
@vectorize([float64(float64, float64)], nopython=True)
def f(x, y):
    return x   y

还可以指定多个函数签名,需注意越通用的类型越要排在后面。

代码语言:javascript复制
import numpy as np
from numba import vectorize, int32, int64, float32, float64

@vectorize([int32(int32, int32),
            int64(int64, int64),
            float32(float32, float32),
            float64(float64, float64)])
def f(x, y):

return x   y

此时如果传入其它类型,程序就会报错。

代码语言:javascript复制
>>> a = np.linspace(0, 1 1j, 6)
>>> f(a, a)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: ufunc 'ufunc' not supported for the input types, and the 
inputs could not be safely coerced to any supported types according
to the casting rule ''safe''
代码语言:javascript复制
你可能会问,为什么不之间用@jit装饰器写一个简单的循环来替代呢?
答案是NumPy UFUNC会自动获得其他功能,如 reduce,accumulate或
广播。

a = np.arange(12).reshape(3, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> f.reduce(a, axis=0) # 按行求和合并
array([12, 15, 18, 21])
>>> f.reduce(a, axis=1) # 按列求和合并
array([ 6, 22, 38])

>>> f.accumulate(a) # 按行积累
array([[ 0,  1,  2,  3],
       [ 4,  6,  8, 10],
       [12, 15, 18, 21]])
>>> f.accumulate(a, axis=1) # 按列积累
array([[ 0,  1,  3,  6],
       [ 4,  9, 15, 22],
       [ 8, 17, 27, 38]])

0 人点赞