计算阵列的期望，提高速度的建议

Question

我有一个 (N0, N1, N2, N3) 矩阵 V 和一个 (N1, N1) 矩阵 M。N1 通常约为 30-50，N0xN1xN2xN3 约为 1 000 000。我想要一个新的 Matrix EV，用于其中 i0, i1, i2, i3 条目由下式给出

np.sum(V[i0, :, i2, i3] * M[i1, :])    

我目前的代码是：

  V_exp = np.tile(V[:, :, :, :, None], (1, 1, 1, 1, N1))
  M_exp = np.tile(M.T[None, :, None, None, :], (N0, 1, N2, N3, 1))
  EV = np.sum(V_exp * M_exp, axis = 1)
  EV = np.rollaxis(EV, 3, 1)

我必须多次执行此操作，这是我代码的绝对瓶颈。我想知道我的代码是否有可能提高速度。我很欣赏建议！

Answer 1

在使用np.newaxis/None 将V 扩展为5D 形状后，对np.einsum 的一次调用将一次性执行所有这些操作，就像这样 -

EV  = np.einsum('ijklm,mj->imkl',V[...,None],M)

因此，我们避免使用任何中间数组以获得内存高效的解决方案。

说明

(1) 涉及扩展维度的起始代码：

V_exp = np.tile(V[:, :, :, :, None], (1, 1, 1, 1, N1))
M_exp = np.tile(M.T[None, :, None, None, :], (N0, 1, N2, N3, 1))
output1 = V_exp * M_exp

按照np.einsum的方式，它会被翻译成：

np.einsum('ijklm,mj->ijklm',V[...,None],M)

请注意，我们使用mj 而不是通常的jm 来对应M 来对应M.T[None, :, None, None, :]。

(2) 接下来，我们有：

EV = np.sum(V_exp * M_exp, axis = 1)

因此，我们沿 axis = 1 求和，因此 einsum 调用需要将 输出字符串说明符从 ->ijklm 更改为 ->iklm。

(3) 最后：

EV = np.rollaxis(EV, 3, 1)

对于rollaxis 的这次移植，我们只需上推 axis=3 到axis=1 的位置，然后下推 axes=1,2 各一个位置向右。因此，输出字符串说明符将从 ->iklm 更改为 ->imkl 给我们：

np.einsum('ijklm,mj->imkl',V[...,None],M)