numpy：高效的大点积

Question

我正在尝试执行大型线性代数计算以将通用协方差矩阵KK_l_obs（形状(NL, NL)）转换为缩小空间中的协方差矩阵映射Kmap_PC（形状(q, q, X, Y)）。

有关如何为每个空间位置构造Kmap_PC 的信息保存在其他数组a、I0 和k_l_th 中。前两个有形状(X, Y)，第三个有(nl, nl)。观察空间和缩小空间之间的转换由 eingenvectors E（形状 (q, nl)）处理。注意NL > nl。

Kmap_PC 的空间元素计算为：

Kmap_PC[..., X, Y] = E.dot(
    KK_l_obs[I0[X, Y]: I0[X, Y] + nl,
             I0[X, Y]: I0[X, Y] + nl] / a_map[X, Y] + \
    k_l_th).dot(E.T)

第一个点积中的位理论上可以直接使用np.einsum 计算，但会占用数百 GB 的内存。我现在正在做的是循环遍历Kmap_PC 的空间索引，这非常慢。我还可以使用 MPI 分配计算（这可能会加速 3-4 倍，因为我有 16 个内核可用）。

我想知道：

(a) 如果我可以更有效地进行计算——也许明确地将其分解为空间元素组；和

(b) 如果我可以改善这些计算的内存开销。

代码 sn-p

import numpy as np
np.random.seed(1)

X = 10
Y = 10
NL = 3000
nl = 1000
q = 7

a_map = 5. * np.random.rand(X, Y)
E = np.random.randn(q, nl)

# construct constant component
m1_ = .05 * np.random.rand(nl, nl)
k_l_th = m1_.dot(m1_)

# construct variable component
m2_ = np.random.rand(NL, NL)
KK_l_obs = m2_.dot(m2_.T)

# where to start in big cov
I0 = np.random.randint(0, NL - nl, (X, Y))

# the slow way
def looping():
    K_PC = np.empty((q, q, X, Y))
    inds = np.ndindex((X, Y))

    for si in inds:
        I0_ = I0[si[0], si[1]]
        K_PC[..., si[0], si[1]] = E.dot(
            KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl] / a_map[si[0], si[1]] + k_l_th).dot(E.T)

    return K_PC

def veccalc():
    nl_ = np.arange(nl)[..., None, None]
    I, J = np.meshgrid(nl_, nl_)

    K_s = KK_l_obs[I0[..., None, None] + J, I0[..., None, None] + I]
    K_s = K_s / a_map[..., None, None] + k_l_th[None, None, ...]
    print(K_s.nbytes)

    K_PC = E @ K_s @ E.T
    K_PC = np.moveaxis(K_PC, [0, 1], [-2, -1])

    return K_PC

Answer 1

调整 #1

在 NumPy 中几乎被忽略的一个非常简单的性能调整是避免使用除法和使用乘法。在处理相同形状的数组时，在处理标量到标量或数组到数组的划分时，这并不明显。但是 NumPy 的隐式广播使得允许在不同形状的数组之间或数组和标量之间广播的划分变得有趣。对于这些情况，我们可以使用与倒数相乘来获得显着的提升。因此，对于所述问题，我们将预先计算 a_map 的倒数，并将其用于乘法而不是除法。

所以，一开始就做：

r_a_map = 1.0/a_map

然后，在嵌套循环中，将其用作：

KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl] * r_a_map[si[0], si[1]]

调整 #2

我们可以在那里使用associative 乘法属性：

A*(B + C) = A*B + A*C

因此，在所有迭代中求和但保持不变的k_l_th 可以在循环之外取出，并在退出嵌套循环后求和。它的有效总和是：E.dot(k_l_th).dot(E.T)。因此，我们会将其添加到 K_PC。

---

定稿和基准测试

使用tweak #1 和tweak#2，我们最终会得到一个修改后的方法，就像这样 -

def original_mod_app():
    r_a_map = 1.0/a_map
    K_PC = np.empty((q, q, X, Y))
    inds = np.ndindex((X, Y))
    for si in inds:
        I0_ = I0[si[0], si[1]]
        K_PC[..., si[0], si[1]] = E.dot(
            KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl] * \
            r_a_map[si[0], si[1]]).dot(E.T)
    return K_PC + E.dot(k_l_th).dot(E.T)[:,:,None,None]

使用与问题中使用的相同示例设置进行运行时测试 -

In [458]: %timeit original_app()
1 loops, best of 3: 1.4 s per loop

In [459]: %timeit original_mod_app()
1 loops, best of 3: 677 ms per loop

In [460]: np.allclose(original_app(), original_mod_app())
Out[460]: True

所以，我们在那里获得了 2x+ 的加速。

Answer 2

在相对适中的机器（4G 内存）上，可以对整个 10x10x1000x1000 空间进行 matmul 计算。

def looping2(n=2):
    ktemp = np.empty((n,n,nl,nl))
    for i,j in np.ndindex(ktemp.shape[:2]):
        I0_ = I0[i, j]
        temp = KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl]
        temp = temp / a_map[i,j] + k_l_th
        ktemp[i,j,...] = temp
    K_PC = E @ ktemp @ E.T      
    return K_PC

K = loop()
k4 = looping2(n=X)
np.allclose(k4, K.transpose(2,3,0,1))  # true

我没有尝试矢量化IO_ 映射。我的重点是推广双点积。

等价的einsum是：

K_PC = np.einsum('ij,...jk,lk->il...', E, ktemp, E) 

这会在 n=7 时产生 ValueError: iterator is too large 错误。

但使用最新版本

K_PC = np.einsum('ij,...jk,lk->il...', E, ktemp, E, optimize='optimal')

适用于完整的 7x7x10x10 输出。

时机不乐观。原始looping 为 2.2 秒，大 matmul（或 einsum）为 3.9 秒。 （我使用original_mod_app 获得相同的 2 倍加速）

============

构造 (10,10,1000,1000) 数组的时间（迭代）：

In [31]: %%timeit 
    ...:     ktemp = np.empty((n,n,nl,nl))
    ...:     for i,j in np.ndindex(ktemp.shape[:2]):
    ...:         I0_ = I0[i, j]
    ...:         temp = KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl]
    ...:         ktemp[i,j,...] = temp
    ...:     
1 loop, best of 3: 749 ms per loop

使用 @ 将其减少到 (10,10,7,7) 的时间（比构造时间长）

In [32]: timeit E @ ktemp @ E.T
1 loop, best of 3: 1.17 s per loop

相同的两个操作的时间，但循环减少

In [33]: %%timeit 
    ...:     ktemp = np.empty((n,n,q,q))
    ...:     for i,j in np.ndindex(ktemp.shape[:2]):
    ...:         I0_ = I0[i, j]
    ...:         temp = KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl]
    ...:         ktemp[i,j,...] = E @ temp @ E.T

1 loop, best of 3: 858 ms per loop

在循环中执行点积可以减少保存到ktemp的子数组的大小，从而弥补计算成本。大数组上的点操作本身比循环更昂贵。即使我们可以“矢量化”KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl]，也无法弥补处理大数组的成本。