Numba矢量化比3d阵列的numpy慢吗?
问题描述:
在下面的代码中,test_func_1大约比test_func_2慢一个数量级。对于此操作是否无法改善或甚至不能匹配numpy性能?Numba矢量化比3d阵列的numpy慢吗?
from numba import guvectorize
import numpy as np
@guvectorize(['void(float64[:,:,:], float64[:], float64[:,:,:])'], '(n,o,p),(n)->(n,o,p)', nopython=True)
def test_func_1(time_series, areas, res):
for i in range(areas.size):
area = areas[i]
adjusted_area = (area/10000.) ** .12 # used to adjust erosion
for k in range(time_series.shape[0]):
res[i, 0, k] = time_series[i, 0, k] * area
res[i, 1, k] = time_series[i, 1, k] * adjusted_area
res[i, 2, k] = time_series[i, 2, k] * area
res[i, 3, k] = time_series[i, 3, k] * adjusted_area
def test_func_2(time_series, areas):
array = np.swapaxes(time_series, 0, 2)
array[:, :2] *= areas
array[:, 2:] *= (areas/10000.) ** .12
return array
dummy = np.float32(np.random.randint(0, 10, (20, 5, 5000)))
areas = np.float32(np.random.randint(0, 10, 20))
test_func_1(dummy, areas)
test_func_2(dummy, areas)
答
由于@JoshAdel在评论中指出的,这里的关键区别是,你numba
版本分配和填充新的阵列,而numpy的是修改,代替原来的。
将适当的.copy()
添加到numpy中会使其稍慢。你也可以使用你的numba版本 - 就我所知不可能使用gufunc而言,但如果你不需要gufunc提供的广播,那么使用常规的jit函数。
def test_func_2(time_series, areas):
array = np.swapaxes(time_series, 0, 2).copy()
array[:, :2] *= areas
array[:, 2:] *= (areas/10000.) ** .12
return array
dummy = np.float32(np.random.randint(0, 10, (20, 5, 5000)))
areas = np.float32(np.random.randint(0, 10, 20))
%timeit test_func_1(dummy, areas)
1.21 ms ± 5.33 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit test_func_2(dummy, areas)
1.77 ms ± 15.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
我使用给定的数据集为'test_func_2(dummy,areas)'获得了大约80微秒的时间。这真的是你的瓶颈吗?或者你只是想学习numba?或者你实际上在处理更大的数据? – Divakar
我得到'test_func_1'大约慢了2倍。当你计算它时,你是否只计算一次(包括jit编译时间)还是计算后续调用,这些调用会缓存编译,并且只是函数的运行时间? – JoshAdel
为了清晰起见,它被简化了。我打算使用类似的功能,最终会在非常大的3D阵列上被称为几万次。 – triphook