我有一个形状为(M,N)的阵列A现在我想做手术
R = (A[:,newaxis,:] * A[newaxis,:,:]).sum(2)
应该产生一个(MxM)数组。现在问题是数组很大,我得到一个内存错误,因为MxMxN数组不适合内存。
完成这项工作的最佳策略是什么? C?地图()?还是有特殊功能吗?
谢谢你, 大卫答案 0 :(得分:7)
我不确定你的数组有多大,但以下是等价的:
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
解决方案更快。完全赞同他指出这一点,但如果他不把它作为一个正式的解决方案发布,我想把它拉出来作为主答案。