numpy：高效的大点产品

Question

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

有关如何Kmap_PC为每个空间位置进行构造的信息保存在其他数组a中I0，和k_l_th。前两个具有形状(X, Y)，第三个具有形状(nl, nl)。观测空间与缩小空间之间的变换由特征向量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分发计算结果（由于我有16个可用内核，因此可以使速度提高3-4倍）。

我很好奇：

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

（b）是否可以改善这些计算的内存开销。

程式码片段

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+ 。