扁平化大数组的 Numpy 平均值比所有轴的平均值慢

Question

运行 Numpy 1.19.2 版，与通过计算已经展平的数组的平均值相比，我在累积数组每个单独轴的平均值方面获得了更好的性能。

shape = (10000,32,32,3)
mat = np.random.random(shape)

# Call this Method A.
%%timeit
mat_means = mat.mean(axis=0).mean(axis=0).mean(axis=0)

14.6 ms ± 167 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

mat_reshaped = mat.reshape(-1,3)

# Call this Method B
%%timeit
mat_means = mat_reshaped.mean(axis=0)

135 ms ± 227 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

这很奇怪，因为多次执行平均值具有与重构阵列上的相同的错误访问模式（可能甚至更糟）。我们也以这种方式进行更多的操作。作为完整性检查，我将数组转换为 FORTRAN 顺序：

mat_reshaped_fortran = mat.reshape(-1,3, order='F')

%%timeit
mat_means = mat_reshaped_fortran.mean(axis=0)

12.2 ms ± 85.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这产生了我预期的性能改进。

对于方法 A，prun给出：

36 function calls in 0.019 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        3    0.018    0.006    0.018    0.006 {method 'reduce' of 'numpy.ufunc' objects}
        1    0.000    0.000    0.019    0.019 {built-in method builtins.exec}
        3    0.000    0.000    0.019    0.006 _methods.py:143(_mean)
        3    0.000    0.000    0.000    0.000 _methods.py:59(_count_reduce_items)
        1    0.000    0.000    0.019    0.019 <string>:1(<module>)
        3    0.000    0.000    0.019    0.006 {method 'mean' of 'numpy.ndarray' objects}
        3    0.000    0.000    0.000    0.000 _asarray.py:86(asanyarray)
        3    0.000    0.000    0.000    0.000 {built-in method numpy.array}
        3    0.000    0.000    0.000    0.000 {built-in method numpy.core._multiarray_umath.normalize_axis_index}
        6    0.000    0.000    0.000    0.000 {built-in method builtins.isinstance}
        6    0.000    0.000    0.000    0.000 {built-in method builtins.issubclass}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

而对于方法 B：

    14 function calls in 0.166 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.166    0.166    0.166    0.166 {method 'reduce' of 'numpy.ufunc' objects}
        1    0.000    0.000    0.166    0.166 {built-in method builtins.exec}
        1    0.000    0.000    0.166    0.166 _methods.py:143(_mean)
        1    0.000    0.000    0.000    0.000 _methods.py:59(_count_reduce_items)
        1    0.000    0.000    0.166    0.166 <string>:1(<module>)
        1    0.000    0.000    0.166    0.166 {method 'mean' of 'numpy.ndarray' objects}
        1    0.000    0.000    0.000    0.000 _asarray.py:86(asanyarray)
        1    0.000    0.000    0.000    0.000 {built-in method numpy.array}
        1    0.000    0.000    0.000    0.000 {built-in method numpy.core._multiarray_umath.normalize_axis_index}
        2    0.000    0.000    0.000    0.000 {built-in method builtins.isinstance}
        2    0.000    0.000    0.000    0.000 {built-in method builtins.issubclass}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

笔记： np.setbufsize(1e7)似乎没有任何影响。

这种性能差异的原因是什么？

Answer 1

我们称您的原始矩阵为mat。mat.shape = (10000,32,32,3)。从视觉上看，这就像有一个由 10,000 * 32x32x3 * 矩形棱柱（我认为它们是乐高积木）组成的“堆栈”。

现在让我们考虑一下您在浮点运算（触发器）方面做了什么：

在方法 A 中，您可以执行mat.mean(axis=0).mean(axis=0).mean(axis=0). 让我们来分解一下：

1. 您取所有 10,000 个乐高积木中每个位置(i,j,k)的平均值。这将返回一个尺寸为32x32x3的乐高积木，其中包含第一组方法。这意味着您执行了 10,000 次加法，平均执行了 1 次除法，其中有 32 32 3 = 3072 次。总共，您执行了30,723,072 次失败。
2. 然后，您再次取每个位置(j,k)的平均值，其中i现在是您当前所在的层（垂直位置）的编号。这将为您提供一张写有32x3平均值的纸。您平均执行了 32 次加法和 1 次除法，其中 32*3 = 96 次。总共执行了3,168 次失败。
3. 最后，取每列k的平均值，其中j现在是您当前所在的行。这将为您提供一个存根，上面写有3 个均值。您已执行 32 次加法，平均执行 1 次除法，其中有 3 次。总共，您完成了99 次失败。

所有这些的总数是 30,723,072 + 3,168 + 99 = 30,726,339flops。

在方法 B 中，您可以这样做 mat_reshaped = mat.reshape(-1,3); mat_means = mat_reshaped.mean(axis=0)。让我们来分解一下：

1. 你重塑了一切，mat一卷尺寸为10,240,000x3的长纸也是如此。您取每列k的平均值，其中j现在是您当前所在的行。这将为您提供一个存根，上面写有3 个均值。您已经执行了 10,240,000 次加法，平均执行了 1 次除法，其中有 3 次。总共，您完成了30,720,003 flops。

所以现在你对自己说：“什么！所有这些工作，只是为了表明较慢的方法实际上工作〜更少〜工作？！”问题是：虽然方法 B 要做的工作较少，但它没有要做的工作少了很多，这意味着仅从失败的角度来看，我们期望事情在运行时方面是相似的。

您还必须考虑方法 B 中重塑数组的大小：具有 10,240,000 行的矩阵非常巨大！对于计算机来说，访问所有这些内容确实很困难/效率低下，而且更多的内存访问意味着更长的运行时间。事实是，在其原始的10,000x32x32x3形状中，矩阵已经被划分为方便的切片，计算机可以更有效地访问：这实际上是处理巨型矩阵时的常用技术Jaime 对类似问题甚至这篇文章的回答：两者都在谈论关于如何将大矩阵分解成更小的切片可以帮助您的程序提高内存效率，从而使其运行得更快。