广播用于元素乘法的3d数组

Question

晚上好,

我需要一些帮助来理解复杂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个元素的求和,以及endproduct:

finalarray:50000x10x10

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

Answer 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