numpy 外积向量化和？

Question

我在 python numpy 中有两个浮点数组：形状为 (60000,) 的

A

和形状为 (60000,784) 的

B

。我尝试了以下操作：

np.einsum('i,ijk->jk',A,B[:,:,np.newaxis]*B[:,np.newaxis,:])

但它会产生内存不足错误，因为在中间步骤中创建形状数组 (60000,784,784) 需要 137GB RAM。

我尝试了以下替代方法：

S=np.zeros((784,784))
for a,b in izip(A,B):
    S+=a*np.outer(b,b)

但是显式的 python sum 太慢了。

是否有一种 numpy 方法可以做到这一点，使用向量化函数来获得快速性能，但在中间步骤中不生成任何 (60000,784,784) 形状数组？

Answer 1

看来你可以简化路径，然后就不需要那么大的中间数组了。

但是不要忘记

optimize=True

，否则会很慢。

import numpy as np

a, b = 60, 784

A = np.random.rand(a)
B = np.random.rand(a, b)

R = np.einsum("i,ijk->jk", A, B[:, :, np.newaxis] * B[:, np.newaxis, :], optimize=True)
R2 = np.einsum("i,ij,ik->jk", A, B, B, optimize=True)

np.allclose(R, R2)  # True