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

用户6021899

发布于 2022-01-10 08:09:09

8580

发布于 2022-01-10 08:09:09

@vectorize 装饰器

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

一个简单的例子：

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

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

@vectorize([float64(float64, float64)], nopython=True)
def f(x, y):
    return x + y

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

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

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

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

你可能会问，为什么不之间用@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]])

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2022-01-06，如有侵权请联系 cloudcommunity@tencent.com 删除

打包