为什么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)

这个方法比你提到的两种方法要快得多。

$ 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 写的，比较通用，因此没有数组版本那么多优化的机会。）