python - 为什么 numpy.apply_along_axis 似乎比 Python 循环慢？

Question

我对 numpy 的 numpy.apply_along_axis() 函数何时会优于简单的 Python 循环感到困惑。例如，考虑一个有很多行的矩阵，您希望计算每一行的总和:

x = np.ones([100000, 3])
sums1 = np.array([np.sum(x[i,:]) for i in range(x.shape[0])])
sums2 = np.apply_along_axis(np.sum, 1, x)

在这里，我什至使用了一个内置的 numpy 函数，np.sum，但是计算 sums1(Python 循环)在计算 时花费了不到 400 毫秒sums2 (apply_along_axis) 需要超过 2000 毫秒(Windows 上的 NumPy 1.6.1)。通过进一步比较，R 的 rowMeans 函数通常可以在不到 20 毫秒的时间内完成此操作(我很确定它正在调用 C 代码)，而类似的 R 函数 apply() 可以在大约 600 毫秒内完成。

Answer 1

np.sum 采用 axis 参数，因此您可以简单地使用

计算总和

sums3 = np.sum(x, axis=1)

这比您提出的 2 种方法快得多。

$ python -m timeit -n 1 -r 1 -s "import numpy as np;x=np.ones([100000,3])" "np.apply_along_axis(np.sum, 1, x)"
1 loops, best of 1: 3.21 sec per loop

$ python -m timeit -n 1 -r 1 -s "import numpy as np;x=np.ones([100000,3])" "np.array([np.sum(x[i,:]) for i in range(x.shape[0])])"
1 loops, best of 1: 712 msec per loop

$ python -m timeit -n 1 -r 1 -s "import numpy as np;x=np.ones([100000,3])" "np.sum(x, axis=1)"
1 loops, best of 1: 1.81 msec per loop

(至于为什么 apply_along_axis 比较慢——我不知道，可能是因为该函数是用纯 Python 编写的，并且比数组版本更通用，因此优化机会更少。)