numpy array：快速将短数组分配给带索引的大数组

Question

我想通过索引将短数组中的值分配给大数组。简单代码如下：

\n

import numpy as np\n\ndef assign_x():\n    a = np.zeros((int(3e6), 20))\n    index_a = np.random.randint(int(3e6), size=(int(3e6), 20))\n    b = np.random.randn(1000, 20)\n    for i in range(20):\n        index_b = np.random.randint(1000, size=int(3e6))\n        a[index_a[:, i], i] = b[index_b, i]\n    return a\n\n%timeit x = assign_x()\n# 2.79 s \xc2\xb1 18.4 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n

\n

我尝试过其他可能相关的方法，例如np.takenumba's jit，但似乎上面是最快的方法。它也可以使用 来加速multiprocessing。我有profile代码，最多的时间是在下面的行，因为它运行了很多次（这里是 20 次）

\n

a[index_a[:, i], i] = b[index_b, i]\n

\n

在使用之前我有机会让它更快吗multiprocessing？

\n

Answer 1

为什么这很慢

这很慢，因为内存访问模式效率很低。事实上，随机访问速度很慢，因为处理器无法预测它们。因此，它会导致昂贵的缓存未命中（如果数组不适合 L1/L2 缓存），而提前预取数据无法避免这种情况。问题是数组太大，无法容纳缓存：index_a每个a数组需要 457 MiB，b需要 156 KiB。因此，b对 L2 缓存的访问通常在具有较高延迟的情况下完成，而对其他两个阵列的访问则在 RAM 中完成。这很慢，因为当前 DDR RAM在典型 PC 上具有 60-100 ns 的巨大延迟。更糟糕的是：这种延迟在不久的将来可能不会小很多：自过去二十年以来，RAM 延迟并没有太大变化。这就是所谓的记忆墙。另请注意，当请求随机位置处的值时，现代处理器通常会从 RAM 中获取通常为 64 字节的完整高速缓存行（导致仅浪费 56/64=87.5% 的带宽）。最后，生成随机数是一个相当昂贵的过程，尤其是大整数，并且np.random.randint可以生成关于目标平台的 32 位或 64 位整数。

如何改善这一点

第一个改进是更喜欢在最连续的维度上进行间接寻址，该维度通常是最后一个维度，因为a[:,i]比 慢a[i,:]。您可以转置数组并交换索引值。然而，Numpy 转置函数只返回一个视图，并没有真正转置内存中的数组。因此当前需要一个明确的副本。这里最好的方法是直接生成数组，以便访问高效（而不是使用昂贵的转置）。请注意，您可以使用简单精度，以便数组可以更好地适应缓存，但代价是精度较低。

下面是一个返回转置数组的示例：

import numpy as np

def assign_x():
    a = np.zeros((20, int(3e6)))
    index_a = np.random.randint(int(3e6), size=(20, int(3e6)))
    b = np.random.randn(20, 1000)
    for i in range(20):
        index_b = np.random.randint(1000, size=int(3e6))
        a[i, index_a[i, :]] = b[i, index_b]
    return a

%timeit x = assign_x()

可以使用Numba进一步改进代码，以便并行运行代码（由于 RAM 延迟，一个核心不应足以使内存饱和，但许多核心可以更好地使用它，因为可以同时完成多个提取）。此外，它可以帮助避免创建大型临时数组。

这是优化 Numba 代码：

import numpy as np
import numba as nb
import random

@nb.njit('float64[:,:]()', parallel=True)
def assign_x():
    a = np.zeros((20, int(3e6)))
    b = np.random.randn(20, 1000)
    for i in nb.prange(20):
        for j in range(3_000_000):
            index_a = random.randint(0, 3_000_000)
            index_b = random.randint(0, 1000)
            a[i, index_a] = b[i, index_b]
    return a

%timeit x = assign_x()

以下是 10 核 Skylake Xeon 处理器上的结果：

Initial code:                  2798 ms
Better memory access pattern:  1741 ms
With Numba:                     318 ms

请注意，并行化最内层循环理论上会更快，因为一行a更有可能适合最后一级缓存。然而，这样做会导致竞争条件，只能通过 Numba（在 CPU 上）尚不可用的原子存储来有效修复该竞争条件。

请注意，最终的代码不能很好地扩展，因为它受内存限制。这是因为 87.5% 的内存吞吐量被浪费了，如前所述。此外，在许多处理器（如所有 Intel 和 AMD-Zen 处理器）上，在这种情况下，写入分配缓存策略强制从内存中读取每个存储的数据。这使得计算效率大大降低，浪费的吞吐量提高到 93.7%……据我所知，Python 中没有办法阻止这种情况。在 C/C++ 中，可以使用低级指令修复写入分配问题。经验法则是避免大阵列上的内存随机访问模式，就像瘟疫一样。