Numba矢量化比使用3d数组的numpy慢？答案

【问题标题】：Numba vectorization slower than numpy with 3d array?Numba矢量化比使用3d数组的numpy慢？
【发布时间】：2017-10-04 16:19:59
【问题描述】：

在下面的代码中，test_func_1 大约比 test_func_2 慢一个数量级。是否无法提高甚至匹配此操作的 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)

【问题讨论】：

对于给定的数据集，test_func_2(dummy, areas) 大约需要 80 u-sec。这真的是你的瓶颈吗？或者你只是想学习 numba？还是您实际上在处理更大的数据？
我知道test_func_1 大约慢了 2 倍。当您计时时，您是只计时一次（包括 jit 编译时间）还是计时后续调用，这些调用将缓存编译并且只是函数的运行时？
为了清楚起见，它被简化了。我打算使用一个类似的函数，最终将在非常大的 3d 数组上调用数万次。
运行 1k 次并取平均值
另外，如果您没有注意到，一个通用的警告词，但 numpy>1.10 中的 np.swapaxes 返回一个视图，因此在示例中，您每次修改 time_series 函数时调用。

标签： python numpy numba

【解决方案1】：

正如@JoshAdel 在 cmets 中指出的那样，这里的主要区别在于您的 numba 版本正在分配和填充新数组，而 numpy 正在就地修改原始数组。

在 numpy 中添加适当的 .copy() 会使我的速度稍微慢一些。你也可以让你的 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)

【讨论】：