我正在尝试执行大型线性代数计算，将通用协方差矩阵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＆gt; 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）如果我可以改善这些计算的内存开销。
代码段
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
答案 0 :(得分：5)
在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]]
我们可以在那里使用乘法的associative属性：
A*(B + C) = A*B + A*C
因此，在所有迭代中求和但保持不变的k_l_th可以在循环之外取出并在离开嵌套循环后求和。它的有效总和将是：E.dot(k_l_th).dot(E.T)。因此，我们会将此添加到K_PC。
使用调整＃1和调整＃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+ 。
答案 1 :(得分：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的子阵列的大小，从而弥补了计算成本。大数组上的点操作本身比你的循环更昂贵。即使我们可以“矢量化”＆＃39; KK_l_obs[I0_ : I0_ + nl, I0_ : I0_ + nl]它无法弥补大数组的成本处理。