晚上好,
我需要一些帮助来理解使用复杂的 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 循环? 谢谢!
对于这样的事情,我会使用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
这应该有效:
(A[:, :, None, None] * B[None, :, :]).sum(axis=1)
但它会炸毁您对产品创建的中间阵列的内存。
该产品具有形状(50000, 2000, 10, 10)
,因此包含 100 亿个元素,对于 64 位浮点值为 80 GB。