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

Question

我想将以下问题并行化。给定一个数组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_samples, dim1, dim2)，且最终和的形状为(dim1, dim2)

我注意到np.einsum对于大型矩阵A来说相当慢。有没有更有效的方法来解决这个问题？我还想尝试使用np.tensordot，但可能情况并非如此

谢谢：-）

Answer 1

In [455]: W = np.arange(1,7).reshape(2,3); A = np.arange(1,13).reshape(3,4)

你的计算：

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]])