在一个太大的数组内部进行乘法和加法

Question

我有一个形状(M,N)的数组A现在我想做手术

R = (A[:,newaxis,:] * A[newaxis,:,:]).sum(2)

它应该产生一个(MxM)数组.现在的问题是数组非常大,我得到一个内存错误,因为MxMxN数组不适合内存.

完成这项工作的最佳策略是什么？C？地图()？还是有特殊功能吗？

谢谢你,大卫

Answer 1

我不确定你的数组有多大,但以下是等价的:

R = np.einsum('ij,kj',A,A)

而且可以更快,并且内存密集度更低:

In [7]: A = np.random.random(size=(500,400))

In [8]: %timeit R = (A[:,np.newaxis,:] * A[np.newaxis,:,:]).sum(2)
1 loops, best of 3: 1.21 s per loop

In [9]: %timeit R = np.einsum('ij,kj',A,A)
10 loops, best of 3: 54 ms per loop

如果我增加的大小A来(500,4000),np.einsum犁通过在约2秒的计算,而原始制剂研我的机器停顿由于它必须创建临时数组的大小.

更新:

正如@Jaime在评论中指出的那样,np.dot(A,A.T)也是问题的等效表述,并且甚至可以比np.einsum解决方案更快.完全赞同他指出这一点,但如果他不将其作为正式解决方案发布,我想把它拉出来作为主要答案.