广播的NumPy算法 - 为什么一种方法的性能更高？

Question

这个问题是我以高效的方式计算Vandermonde矩阵的答案 .

这是设置:

x = np.arange(5000)  # an integer array
N = 4

现在,我将以两种不同的方式计算Vandermonde矩阵:

m1 = (x ** np.arange(N)[:, None]).T

和,

m2 = x[:, None] ** np.arange(N)

完整性检查:

np.array_equal(m1, m2)
True

这些方法是相同的,但它们的性能不是:

%timeit m1 = (x ** np.arange(N)[:, None]).T
42.7 µs ± 271 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit m2 = x[:, None] ** np.arange(N)
150 µs ± 995 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

因此,第一种方法,尽管最后要求换位,仍然比第二种方法快3倍.

唯一的区别是,在第一种情况下,较小的阵列是广播的,而在第二种情况下,它是较大的.

所以,通过对numpy如何工作的相当不错的理解,我可以猜测答案将涉及缓存.第一种方法比第二种方法更加缓存友好.但是,我想要一个比我更有经验的人的官方消息.

在时间上形成鲜明对比的原因是什么？

Answer 1

虽然我担心我的结论不会比你的更基本（“可能是缓存”），但我相信我可以通过一组更本地化的测试来帮助我们集中注意力。

\n\n

考虑您的示例问题：

\n\n

M,N = 5000,4\nx1 = np.arange(M)\ny1 = np.arange(N)[:,None]\nx2 = np.arange(M)[:,None]\ny2 = np.arange(N)\nx1_bc,y1_bc = np.broadcast_arrays(x1,y1)\nx2_bc,y2_bc = np.broadcast_arrays(x2,y2)\nx1_cont,y1_cont,x2_cont,y2_cont = map(np.ascontiguousarray,\n                                      [x1_bc,y1_bc,x2_bc,y2_bc])\n

\n\n

正如你所看到的，我定义了一堆数组来进行比较。x1、y1和x2、y2分别对应于您的原始测试用例。??_bc对应于这些数组的显式广播版本。它们与原始数据共享数据，但它们具有明确的 0 步幅以获得适当的形状。最后，??_cont是这些广播数组的连续版本，就像用np.tile.

\n\n

因此x1_bc、y1_bc、x1_cont和 都y1_cont具有 shape (4, 5000)，但前两者的步长为零，而后两者是连续的数组。出于所有意图和目的，利用任何这些相应数组对的能力应该给我们相同的连续结果（正如 hpaulj 在评论中指出的那样，转置本身本质上是免费的，所以我将忽略最外面的转置在下面的）。

\n\n

以下是与您的原始支票相对应的时间：

\n\n

In [143]: %timeit x1 ** y1\n     ...: %timeit x2 ** y2\n     ...: \n52.2 \xc2\xb5s \xc2\xb1 707 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n96 \xc2\xb5s \xc2\xb1 858 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n

\n\n

以下是显式广播数组的时序：

\n\n

In [144]: %timeit x1_bc ** y1_bc\n     ...: %timeit x2_bc ** y2_bc\n     ...: \n54.1 \xc2\xb5s \xc2\xb1 906 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n99.1 \xc2\xb5s \xc2\xb1 1.51 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each\n

\n\n

一样。这告诉我，差异并不是由于从索引表达式到广播数组的转换造成的。这基本上是预料之中的，但检查一下总没有坏处。

\n\n

最后，连续数组：

\n\n

In [146]: %timeit x1_cont ** y1_cont\n     ...: %timeit x2_cont ** y2_cont\n     ...: \n38.9 \xc2\xb5s \xc2\xb1 529 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n45.6 \xc2\xb5s \xc2\xb1 390 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n

\n\n

很大一部分差异消失了！

\n\n

那么我为什么要检查这个呢？有一个一般的经验法则，如果您在 python 中沿大尾随维度使用矢量化操作，则可以使用 CPU 缓存。更具体地说，对于行主（“C 顺序”）数组，尾随维度是连续的，而对于列主（“fortran 顺序”）数组，前导维度是连续的。对于足够大的维度arr.sum(axis=-1)应该比arr.sum(axis=0)行主 numpy 数组更快，给出或采取一些细则。

\n\n

这里发生的情况是，两个维度（大小分别为 4 和 5000）之间存在巨大差异，但两种转置情况之间的巨大性能不对称仅发生在广播情况下。无可否认，我的印象是广播使用 0 步幅来构建适当大小的视图。这些 0 步长意味着在更快的情况下，长x数组的内存访问如下所示：

\n\n

[mem0,mem1,mem2,...,mem4999, mem0,mem1,mem2,...,mem4999, ...] # and so on\n

\n\n

其中mem*仅表示位于 RAM 中某处float64的值x。将此与我们使用 shape 的较慢情况进行比较(5000,4)：

\n\n

[mem0,mem0,mem0,mem0, mem1,mem1,mem1,mem1, mem2,mem2,mem2,mem2, ...]\n

\n\n

我天真的想法是，使用前者可以让 CPU 一次缓存更大的单个值块x，因此性能非常好。在后一种情况下，0 跨步使 CPU 在同一内存地址上跳跃x四次，总共执行 5000 次。我认为有理由相信此设置不利于缓存，从而导致整体性能不佳。这也与以下事实相一致：连续的情况不会显示出这种性能影响：CPU 必须处理所有 5000*4 个唯一float64值，并且缓存可能不会受到这些奇怪的读取的阻碍。

\n