NumPy tensordot MemoryError

Question

我有两个矩阵 - A是3033x3033，X是3033x20。我正在运行以下行（如the answer to another question I asked中所述）：

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简单地停下来说＆＃34; MemoryError＆＃34;。如何通过更改Python中的某些设置或以更节省内存的方式执行此操作来解决这个问题？

Answer 1

这是一个在没有任何for循环且没有任何大型临时数组的情况下进行计算的函数。请参阅related question获取更长的答案，并附上测试脚本。

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

我用你的size数组分析了代码：

顺便说一句，我喜欢sp.einsum。通过删除for循环加速数组操作是一个很好的起点。你可以通过一次调用sp.einsum来做很多事情。

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

Answer 2

如果用

替换最后一行

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

你避免创建我认为是你主要问题的A.reshape(n, n, -1) * c（3301乘3301乘20）中间件。

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

你可以进一步避免创建c，但我无法立即看到如何做到这一点。这是根据矩阵指标的总和以及查看它简化为什么来编写整个事物的情况。

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])