python - numpy einsum 的替代品

Question

当我计算具有 N 行和 n 列的矩阵 X 的三阶矩时，我通常使用 einsum:

M3 = sp.einsum('ij,ik,il->jkl',X,X,X) /N

这通常工作正常，但现在我使用更大的值，即 n = 120 和 N = 100000，并且 einsum 返回以下错误:

ValueError: iterator is too large

做 3 个嵌套循环的替代方案是不可行的，所以我想知道是否有任何替代方案。

Answer 1

请注意，计算它至少需要执行 ~n³ × N = 1730 亿次操作(不考虑对称性)，因此它会很慢，除非 numpy 可以访问 GPU 或其他东西。在配备 ~3 GHz CPU 的现代计算机上，假设没有 SIMD/并行加速，整个计算预计需要大约 60 秒才能完成。

---

为了测试，让我们从 N = 1000 开始。我们将使用它来检查正确性和性能:

#!/usr/bin/env python3

import numpy
import time

numpy.random.seed(0)

n = 120
N = 1000
X = numpy.random.random((N, n))

start_time = time.time()

M3 = numpy.einsum('ij,ik,il->jkl', X, X, X)

end_time = time.time()

print('check:', M3[2,4,6], '= 125.401852515?')
print('check:', M3[4,2,6], '= 125.401852515?')
print('check:', M3[6,4,2], '= 125.401852515?')
print('check:', numpy.sum(M3), '= 218028826.631?')
print('total time =', end_time - start_time)

这大约需要 8 秒。这是基线。

让我们从 3 个嵌套循环开始:

M3 = numpy.zeros((n, n, n))
for j in range(n):
    for k in range(n):
        for l in range(n):
            M3[j,k,l] = numpy.sum(X[:,j] * X[:,k] * X[:,l])
# ~27 seconds

这大约需要半分钟，不好!一个原因是因为这实际上是四个嵌套循环:numpy.sum 也可以认为是一个循环。

我们注意到可以将总和转换为点积以移除第 4 个循环:

M3 = numpy.zeros((n, n, n))
for j in range(n):
    for k in range(n):
        for l in range(n):
            M3[j,k,l] = X[:,j] * X[:,k] @ X[:,l]
# 14 seconds

现在好多了，但仍然很慢。但是我们注意到点积可以变成矩阵乘法以去除一个循环:

M3 = numpy.zeros((n, n, n))
for j in range(n):
    for k in range(n):
        M3[j,k] = X[:,j] * X[:,k] @ X
# ~0.5 seconds

嗯？现在这比 einsum 还要高效!我们还可以检查答案是否确实正确。

我们可以更进一步吗？是的!我们可以通过以下方式消除 k 循环:

M3 = numpy.zeros((n, n, n))
for j in range(n):
    Y = numpy.repeat(X[:,j], n).reshape((N, n))
    M3[j] = (Y * X).T @ X
# ~0.3 seconds

我们还可以使用广播(即 a * [b,c] == [a*b, a*c] 用于 X 的每一行)来避免执行 numpy。重复(感谢@Divakar):

M3 = numpy.zeros((n, n, n))
for j in range(n):
    Y = X[:,j].reshape((N, 1))
    ## or, equivalently: 
    # Y = X[:, numpy.newaxis, j]
    M3[j] = (Y * X).T @ X
# ~0.16 seconds

如果我们将其缩放到 N = 100000，则程序预计需要 16 秒，这在理论限制内，因此消除 j 可能不会有太大帮助(但这可能会使代码真的很难理解)。我们可以接受这作为最终解决方案。

---

注意:如果您使用的是 Python 2，a @b 等同于 a.dot(b)。