Python 计算一批数组的数组矩阵乘法的有效方法

我想将以下问题并行化。给定一个数组

w

（dim1，）和一个矩阵

a

（dim1，dim2），我希望

a

的每一行乘以

w

的对应元素。 那是很琐碎的

但是，我想对一组数组

w

这样做，并最终对结果求和。因此，为了避免for循环，我使用shape

（n_samples，dim1）

创建了矩阵

W

，并以以下方式使用

np.einsum

函数：

x = np.einsum('ji, ik -> jik', W, A))
r = x.sum(axis=0)

其中，

x

的形状为

（n_样本，dim1，dim2）

，且最终总和的形状为

（dim1，dim2）

我注意到，

np.einsum

对于大型矩阵

a

来说相当慢。有没有更有效的方法来解决这个问题？我还想尝试使用

np.tensordot

，但可能不是这样

谢谢：-）

你的计算：

In [463]: x = np.einsum('ji, ik -> jik', W, A)
     ...: r = x.sum(axis=0)
In [464]: r
Out[464]: 
array([[  5,  10,  15,  20],
       [ 35,  42,  49,  56],
       [ 81,  90,  99, 108]])

如注释中所述，

einsum

可以在

j

上执行求和：

In [465]: np.einsum('ji, ik -> ik', W, A)
Out[465]: 
array([[  5,  10,  15,  20],
       [ 35,  42,  49,  56],
       [ 81,  90,  99, 108]])

由于

j

只出现在

A

中，我们可以先对

A

求和：

In [466]: np.sum(W,axis=0)[:,None]*A
Out[466]: 
array([[  5,  10,  15,  20],
       [ 35,  42,  49,  56],
       [ 81,  90,  99, 108]])

这不涉及乘积和，矩阵乘法也不涉及

或者在乘法后求和：

In [475]: (W[:,:,None]*A).sum(axis=0)
Out[475]: 
array([[  5,  10,  15,  20],
       [ 35,  42,  49,  56],
       [ 81,  90,  99, 108]])

---

您可以在输出中省略

j

下标，而不需要将数组求和作为额外步骤。Tensordot实际上只是einsum的一个推广。如果您想要明确，应该使用einsum。大多数优化都是在numpy的引擎盖下完成的：）

tensordot

对其输入进行重塑和转换，将问题简化为一个

np.dot

，之后进行一些进一步的按摩<代码>np.einsum更为通用

np.matmul/@

是批处理矩阵乘法函数。

In [475]: (W[:,:,None]*A).sum(axis=0)
Out[475]: 
array([[  5,  10,  15,  20],
       [ 35,  42,  49,  56],
       [ 81,  90,  99, 108]])