NumPy tensordot 内存错误。

Question

我有两个矩阵——A是3033x3033，X是3033x20。我正在运行以下代码（如我另一个问题的答案中建议的）：

n, d = X.shape
c = X.reshape(n, -1, d) - X.reshape(-1, n, d)
return np.tensordot(A.reshape(n, n, -1) * c, c, axes=[(0,1),(0,1)])

在最后一行，Python会直接停止并显示“MemoryError”错误。我该如何解决这个问题？可以通过更改Python设置或以更节省内存的方式执行这个操作来解决。

- Andrew Latham

---

你期望结果的 .shape 是什么？ - farenorth

另外，在代码中操作A和X之前，您是否声明了它们的形状？n，d是什么？您应该简化您的代码，以便我们可以更轻松地帮助您（即帮助我帮助您）。将最终行分成多个部分，例如：in1 = A.reshape(n, n, -1) * c，然后执行np.tensordot(in1, c, axes=[(0,1), (0,1)])。我猜您的问题出现在第二部分中。然后，您只需说明in1和c的尺寸，而不是强迫我弄清楚发生了什么。 - farenorth

Answer 1

这里有一个函数，它可以在没有任何for循环和大型临时数组的情况下进行计算。详见相关问题，里面有更长的答案和完整的测试脚本。

def fbest(A, X):
  ""
  KA_best = np.tensordot(A.sum(1)[:,None] * X, X, axes=[(0,), (0,)])
  KA_best += np.tensordot(A.sum(0)[:,None] * X, X, axes=[(0,), (0,)])
  KA_best -= np.tensordot(np.dot(A, X), X, axes=[(0,), (0,)])
  KA_best -= np.tensordot(X, np.dot(A, X), axes=[(0,), (0,)])
  return KA_best

我使用您提供的大小数组对代码进行了剖析：

顺便说一下，我很喜欢sp.einsum。它是加速数组操作的绝佳起点，可以通过删除for循环实现此目的。您只需调用一次sp.einsum就可以完成许多工作。

np.tensordot的优点在于它可以链接到您安装的任何快速数值库（即MKL）。因此，当您安装了正确的库时，tensordot将以更快的速度和并行运行。

- Will Martin

Answer 2

如果您将最后一行替换为：

return np.einsum('ij,ijk,ijl->kl',A,c,c)

我认为你的主要问题在于避免创建 A.reshape(n, n, -1) * c (3301乘以3301乘以20) 这个中间过程。

我的印象是，我提供的版本可能会更慢（对于不会耗尽内存的情况），但我没有严格地计时。

可能你可以进一步避免创建 c，但我暂时看不出如何做到。这将是按矩阵索引和总结的形式编写整个内容，并看看它简化成了什么。

- DavidW

Answer 3

您可以使用两层嵌套循环格式来迭代X的最后一个维度。现在，该最后一个维度是20，因此希望它仍然足够高效，并且更重要的是留下最小的内存占用。以下是实现方式 -

n, d = X.shape
c = X.reshape(n, -1, d) - X.reshape(-1, n, d)
out = np.empty((d,d))  # d is a small number: 20
for i in range(d):
    for j in range(d):
        out[i,j] = (A*c[:,:,i]*(c[:,:,j])).sum()

 return out

您可以将最后一行替换为np.einsum -

out[i,j] = np.einsum('ij->',A*c[:,:,i]*c[:,:,j])    

- Divakar