9

我对 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 毫秒,而计算sums2apply_along_axis)需要超过 2000 毫秒(Windows 上的 NumPy 1.6.1)。通过进一步的比较,R 的 rowMeans 函数通常可以在不到 20 毫秒的时间内完成(我很确定它正在调用 C 代码),而类似的 R 函数apply()可以在大约 600 毫秒内完成。

4

1 回答 1

10

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 编写的,并且比数组版本更通用,因此优化机会更少。)

于 2011-12-28T07:12:32.347 回答