广播 3d 数组以进行元素乘法答案

【问题标题】：Broadcasting 3d arrays for elementwise multiplication广播 3d 数组以进行元素乘法
【发布时间】：2018-09-06 18:47:36
【问题描述】：

晚上好，

我需要一些帮助来理解复杂的 numpy 数组的高级广播。

我有：

阵列 A：50000x2000

阵列 B：2000x10x10

用for循环实现：

for k in range(50000):
    temp = A[k,:].reshape(2000,1,1)
    finalarray[k,:,:]=np.sum ( B*temp , axis=0)

我想要一个带有 2000 个元素的轴的元素乘法和求和，最终产品：

最终数组：50000x10x10

是否可以避免 for 循环？谢谢！

【问题讨论】：

如果 B 是 (2000,100)，您可以使用 np.dot。

标签： python numpy broadcasting

【解决方案1】：

对于这样的事情，我会使用np.einsum，这样可以很容易地根据你想要的索引操作写下你想要发生的事情：

fast = np.einsum('ij,jkl->ikl', A, B)

这给了我相同的结果（丢弃 50000->500 以便快速完成循环）：

A = np.random.random((500, 2000))
B = np.random.random((2000, 10, 10))
finalarray = np.zeros((500, 10, 10))
for k in range(500):
    temp = A[k,:].reshape(2000,1,1)
    finalarray[k,:,:]=np.sum ( B*temp , axis=0)

fast = np.einsum('ij,jkl->ikl', A, B)

给我

In [81]: (finalarray == fast).all()
Out[81]: True

即使在 50000 的情况下也有合理的性能：

In [88]: %time fast = np.einsum('ij,jkl->ikl', A, B)
Wall time: 4.93 s

In [89]: fast.shape

Out[89]: (50000, 10, 10)

或者，在这种情况下，您可以使用tensordot：

faster = np.tensordot(A, B, axes=1)

这将快几倍（以不那么通用为代价）：

In [29]: A = np.random.random((50000, 2000))

In [30]: B = np.random.random((2000, 10, 10))

In [31]: %time fast = np.einsum('ij,jkl->ikl', A, B)
Wall time: 5.08 s

In [32]: %time faster = np.tensordot(A, B, axes=1)
Wall time: 504 ms

In [33]: np.allclose(fast, faster)
Out[33]: True

我不得不在这里使用allclose，因为这些值最终会略有不同：

In [34]: abs(fast - faster).max()
Out[34]: 2.7853275241795927e-12

【讨论】：

另一种但比 tensordot 慢的是result = A @ np.transpose(B, (1,0,2))。在这种情况下，它与 einsum 解决方案的速度大致相同。

【解决方案2】：

这应该可行：

(A[:, :, None, None] * B[None, :, :]).sum(axis=1)

但它会炸毁你对产品创建的中间数组的记忆。

产品的形状为(50000, 2000, 10, 10)，因此包含100亿个元素，对于64位浮点值来说是80 GB。

【讨论】：