为什么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，然而计算Python循环中的sums1仅需不到400毫秒，而计算sums2（apply_along_axis）需要超过2000毫秒（在Windows上使用NumPy 1.6.1）。作为进一步的比较，R的rowMeans函数通常可以在不到20毫秒的时间内完成此操作（我很确定它正在调用C代码），而类似的R函数apply()则可以在约600毫秒内完成。

- Abiel

---

不幸的是，沿轴应用似乎只适用于非速度相关任务。 - Wizard

Answer 1

< p >np.sum函数有一个axis参数，因此你可以直接使用它来计算总和。

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

这比你提出的两种方法快得多。

$ 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编写的，而且更加通用，因此优化机会更少，比数组版本要慢。

- kennytm