Einsum 比 n 模式张量矩阵乘积的显式 Numpy 实现慢

Question

我正在尝试使用 Numpy 在 Python 中高效地实现 n 模式张量矩阵乘积（由 Kolda 和 Bader 定义：https://www.sandia.gov/~tgkolda/pubs/pubfiles/SAND2007-6702.pdf）。该操作有效地归结为（对于矩阵 U、张量 X 和轴/模式 k）：

1. 通过折叠所有其他轴，从 X 中提取沿轴 k 的所有向量。

2. 使用标准矩阵乘法将左侧的这些向量乘以 U。

3. 使用相同的形状再次将向量插入输出张量，除了 X.shape[k]，它现在等于 U.shape[0]（最初，X.shape[k] 必须是等于 U.shape[1]，作为矩阵乘法的结果）。

一段时间以来，我一直在使用显式实现，它分别执行所有这些步骤：

1. 转置张量以将轴 k 带到前面（在我的完整代码中，我添加了一个例外情况，以防 k == X.ndim - 1，在这种情况下，将其保留在那里并转置所有未来的操作会更快，或者至少在我的应用程序中，但这与此处无关）。

2. 重塑张量以折叠所有其他轴。

3. 计算矩阵乘法。

4. 重塑张量以重建所有其他轴。

5. 将张量转回原来的顺序。

我认为这个实现会创建很多不必要的（大）数组，所以一旦我发现 np.einsum，我认为这会大大加快速度。但是使用下面的代码我得到了更糟糕的结果：

import numpy as np
from time import time

def mode_k_product(U, X, mode):
    transposition_order = list(range(X.ndim))
    transposition_order[mode] = 0
    transposition_order[0] = mode
    Y = np.transpose(X, transposition_order)
    transposed_ranks = list(Y.shape)
    Y = np.reshape(Y, (Y.shape[0], -1))
    Y = U @ Y
    transposed_ranks[0] = Y.shape[0]
    Y = np.reshape(Y, transposed_ranks)
    Y = np.transpose(Y, transposition_order)
    return Y

def einsum_product(U, X, mode):
    axes1 = list(range(X.ndim))
    axes1[mode] = X.ndim + 1
    axes2 = list(range(X.ndim))
    axes2[mode] = X.ndim
    return np.einsum(U, [X.ndim, X.ndim + 1], X, axes1, axes2, optimize=True)

def test_correctness():
    A = np.random.rand(3, 4, 5)
    for i in range(3):
        B = np.random.rand(6, A.shape[i])
        X = mode_k_product(B, A, i)
        Y = einsum_product(B, A, i)
        print(np.allclose(X, Y))

def test_time(method, amount):
    U = np.random.rand(256, 512)
    X = np.random.rand(512, 512, 256)
    start = time()
    for i in range(amount):
        method(U, X, 1)
    return (time() - start)/amount

def test_times():
    print("Explicit:", test_time(mode_k_product, 10))
    print("Einsum:", test_time(einsum_product, 10))

test_correctness()
test_times()

适合我的时间：

显式：3.9450525522232054

Einsum：15.873924326896667

这是正常的还是我做错了什么？我知道在某些情况下存储中间结果可以降低复杂性（例如链式矩阵乘法），但是在这种情况下，我想不出任何重复的计算。矩阵乘法是否如此优化以至于它消除了不转置的好处（从技术上讲，它的复杂性较低）？

Answer 1

我比较熟悉使用einsum的下标风格，所以算出这些等价：

In [194]: np.allclose(np.einsum('ij,jkl->ikl',B0,A), einsum_product(B0,A,0))          
Out[194]: True
In [195]: np.allclose(np.einsum('ij,kjl->kil',B1,A), einsum_product(B1,A,1))          
Out[195]: True
In [196]: np.allclose(np.einsum('ij,klj->kli',B2,A), einsum_product(B2,A,2))          
Out[196]: True

使用mode 参数，您在einsum_product 中的方法可能是最好的。但是等价性可以帮助我更好地可视化计算，并且可能对其他人有所帮助。

时间应该基本相同。 einsum_product 中有一个额外的设置时间，在更大的维度上应该会消失。

Answer 2

更新 Numpy 后，Einsum 仅比显式方法稍慢，无论是否使用多线程（请参阅我的问题的 cmets）。