Python Numpy - numpy axis 性能

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函数，沿列计算会更快。

所以我的观点是为什么沿axis = 0（沿列计算）的速度比沿axis = 1（沿行计算）的速度更快。

例如

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

在行优先（row major）排序中，[1,2,3]和[4,5,6]，[7,8,9]分别被分配到相邻的内存中。

因此，沿着axis = 1计算速度应该比axis = 0更快。然而，使用numpy sum函数时，沿列（axis = 0）计算速度更快。

您如何解释这一点呢？

谢谢！

- 신대규

---

a[0,:] 只对第一行进行求和。a.sum(axis=0) 沿着行方向对整个矩阵进行求和。这两个操作计算的并不相同。 - cs95

然而，如果使用numpy内置函数，则结果会有所不同。实际上，您始终使用“numpy内置函数”！@COLDSPEED完全正确：您正在计算不同的内容。 - AGN Gazer

a.sum(axis=0) 是按列计算的结果。而沿行计算的 "print(a[0,:].sum(), a[1,:].sum(), a[2,:].sum())" 的结果是 [3 12 21]，而 a.sum(axis=0) 的结果是 [9, 12, 15]。 - 신대규

@신대규 如果你想要相同的输出，使用 axis=1 而不是 0。此外，一个比另一个慢的原因是由于缓存未命中和局部性差。 - cs95

@cᴏʟᴅsᴘᴇᴇᴅ 谢谢您的回复。 这对我来说很困惑。我也在Linux系统上进行了测试。a = np.random.rand(5000, 5000)%timeit a.sum(axis=0) 100次循环，3次取最佳结果：每个循环17毫秒 %timeit a.sum(axis=1) 100次循环，3次取最佳结果：每个循环15.5毫秒 虽然差别不大，但是有意义。但是在Windows系统上。 %timeit a.sum(axis=0) 7次循环，100次取平均值：平均每个循环17毫秒，标准偏差为122微秒 %timeit a.sum(axis=1) 7次循环，10次取平均值：平均每个循环29.7毫秒，标准偏差为153微秒。 看起来好像受操作系统或numpy版本的影响。 - 신대규

Answer 1

你不是在计算同一件事情。

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

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

第二个和第三个命令是对2D数组的全部内容进行求和，但是沿着某个轴进行。这样，您不会得到单一的结果（就像前两个命令一样），而是得到一个包含所有求和结果的1D数组。

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])

- cs95

---

值得注意的是，如果您执行 c = np.random.randn(5000, 5000) 和 f = np.asfortranarray(c)，则在0轴和1轴上进行求和时会看到非常明显的时间差异。 - Brad Solomon