Python Numpy - numpy 轴性能

Question

默认情况下，numpy 是行优先的。因此，我自然而然地接受了以下结果。

a = np.random.rand(5000, 5000)

%timeit a[0,:].sum()
3.57 µs ± 13.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit a[:,0].sum()
38.8 µs ± 8.19 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

因为是行主序，自然是一个[0,:]计算得更快。但是，如果使用 numpy sum 函数，结果就不同了。

%timeit a.sum(axis=0)
16.9 ms ± 13.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit a.sum(axis=1)
29.5 ms ± 90.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

如果使用 numpy sum 函数，沿着列计算它会更快。

所以我的观点是为什么沿轴 = 0(沿列计算)的速度比沿轴 = 1(沿行)快。

例如

a = np.array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]], order='C')

按行主顺序，[1,2,3]和[4,5,6]，[7,8,9]分别分配到相邻内存。

因此，沿 axis = 1 计算的速度应该比 axis = 0 快。但是，当使用 numpy sum 函数时，沿列(轴 = 0)计算速度更快。

你怎么解释这个？

谢谢

Answer 1

你计算的不是同一件事。

前两个命令只计算整个数组中的一行/一列。

a[0, :].sum().shape   # sums just the first row only
()

后两个命令对二维数组的全部内容求和，但沿特定轴。这样，您不会得到单个结果(如前两个命令)，而是一维总和数组。

a.sum(axis=0).shape   # computes the row-wise sum for each column
(5000,)

总而言之，两组命令做的事情不同。

---

a
array([[1, 6, 9, 1, 6],
       [5, 6, 9, 1, 3],
       [5, 0, 3, 5, 7],
       [2, 8, 3, 8, 6],
       [3, 4, 8, 5, 0]])

a[0, :]
array([1, 6, 9, 1, 6])

a[0, :].sum()
23

a.sum(axis=0)
array([16, 24, 32, 20, 22])