python - 扁平大阵列的 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) .让我们分解一下:

您取所有 10,000 个乐高积木中每个位置 (i,j,k) 的平均值。这给了你一个 32x32x3 大小的乐高积木，它现在包含第一组装置。这意味着您已经执行了 10,000 次加法和平均 1 次除法，其中有 32323 = 3072。总共，您已经完成了 30,723,072 次翻牌 .

然后，您再次取平均值，这次是每个位置 (j,k)，其中 i 现在是您当前所在的层数(垂直位置)。这会给你一张纸，上面写着 32x3 的意思。您已经执行了 32 次加法和 1 次除法，其中有 32*3 = 96。总共，您完成了 3,168 次翻牌 .

最后，取每列 k 的平均值，其中 j 现在是您当前所在的行。这给你一个 stub ，上面写着 3 种手段。您已经执行了 32 次加法和 1 次除法，其中有 3 次。总共，您完成了 99 次翻牌 .

所有这些的总和是 30,723,072 + 3,168 + 99 = 30,726,339翻牌。
在方法 B 中，您执行 mat_reshaped = mat.reshape(-1,3); mat_means = mat_reshaped.mean(axis=0) .让我们分解一下:

你 reshape 了一切，所以mat是一长卷纸，大小为 10,240,000x3。您取每列 k 的平均值，其中 j 现在是您当前所在的行。这给你一个 stub ，上面写着 3 种手段。您已经执行了 10,240,000 次加法和 1 次除法，其中有 3 次。总共，您完成了 30,720,003 flops .

所以现在你对自己说:“什么!所有这些工作，只是为了表明较慢的方法确实 ~less~ 工作？!”问题是:虽然方法 B 要做的工作更少，但它没有要做的工作要少得多，这意味着仅从失败的角度来看，我们希望事情在运行时方面是相似的。
您还必须考虑方法 B 中重构数组的大小:10,240,000 行的矩阵是巨大的!!!计算机访问所有这些真的很难/效率低下，更多的内存访问意味着更长的运行时间。事实是，在最初的 10,000x32x32x3 形状中，矩阵已经被划分为方便的切片，计算机可以更有效地访问:这实际上是处理巨型矩阵时的常用技术 Jaime's response to a similar question甚至this article :两者都谈到如何将大矩阵分解成更小的切片有助于您的程序提高内存效率，从而使其运行得更快。