Matrix/Tensor 三重产品？

Question

我正在研究的算法需要在几个地方计算一种矩阵三元组积。

该运算采用三个具有相同维度的方阵，并生成一个 3 索引张量。标记操作数 A、B 和 C，结果的第 (i,j,k) 个元素是

X[i,j,k] = \sum_a A[i,a] B[a,j] C[k,a]

在 numpy 中，您可以使用 einsum('ia,aj,ka->ijk', A, B, C).

来计算

问题：

这个操作有标准名称吗？
我可以通过一次 BLAS 调用来计算吗？
是否有任何其他高度优化的数值 C/Fortran 库可以计算这种类型的表达式？

Answer 1

设 nxn 为矩阵大小。在 Matlab 中，您可以

将 A 和 C 分组为 n^2xn 矩阵 AC，这样 AC 的行对应所有组合A 和 C.
Post-将 AC 乘以 B。这给出了期望的结果，只是形状不同。
重塑和置换维度以获得所需形式的结果。

代码：

AC = reshape(bsxfun(@times, permute(A, [1 3 2]), permute(C, [3 1 2])), n^2, n); % // 1
X = permute(reshape((AC*B).', n, n, n), [2 1 3]);                               %'// 2, 3

使用基于逐字循环的方法进行检查：

%// Example data:
n = 3;
A = rand(n,n);
B = rand(n,n);
C = rand(n,n);

%// Proposed approach:
AC = reshape(bsxfun(@times, permute(A, [1 3 2]), permute(C, [3 1 2])), n^2, n);
X = permute(reshape((AC*B).', n, n, n), [2 1 3]); %'

%// Loop-based approach:
Xloop = NaN(n,n,n); %// initiallize
for ii = 1:n
    for jj = 1:n
        for kk = 1:n
            Xloop(ii,jj,kk) = sum(A(ii,:).*B(:,jj).'.*C(kk,:)); %'
        end
    end
end

%// Compute maximum relative difference:
max(max(max(abs(X./Xloop-1))))

ans =
    2.2204e-16

最大相对差为eps量级，因此结果在数值精度内是正确的。

Answer 2

简介和解决方案代码

np.einsum, is really hard to beat, but in rare cases, you can still beat it, if you can bring in matrix-multiplication into the computations. After few trials, it seems you can bring in matrix-multiplication with np.dot性能超越np.einsum('ia,aj,ka->ijk', A, B, C).

基本思想是我们将“all einsum”操作分解为 np.einsum 和 np.dot 的组合，如下所示：

A:[i,a] 和 B:[a,j] 的求和是用 np.einsum 完成的，得到 3D array:[i,j,a]。
然后将此 3D 数组重塑为 2D array:[i*j,a]，第三个数组 C[k,a] 转置为 [a,k]，目的是在这两个数组之间执行 matrix-multiplication ，给我们 [i*j,k] 作为矩阵乘积，因为我们在那里失去了索引 [a]。
最终输出的产品被重塑为 3D array:[i,j,k]。

这是目前讨论的第一个版本的实现 -

import numpy as np

def tensor_prod_v1(A,B,C):   # First version of proposed method
    # Shape parameters
    m,d = A.shape
    n = B.shape[1]
    p = C.shape[0]
    
    # Calculate \sum_a A[i,a] B[a,j] to get a 3D array with indices as (i,j,a)
    AB = np.einsum('ia,aj->ija', A, B)
    
    # Calculate entire summation losing a-ith index & reshaping to desired shape
    return np.dot(AB.reshape(m*n,d),C.T).reshape(m,n,p)

由于我们对所有三个输入数组的 a-th 索引求和，因此可以使用三种不同的方法对第 a 个索引求和。前面列出的代码用于 (A,B)。因此，我们也可以有 (A,C) 和 (B,C) 给我们另外两个变体，如下所列：

def tensor_prod_v2(A,B,C):
    # Shape parameters
    m,d = A.shape
    n = B.shape[1]
    p = C.shape[0]
    
    # Calculate \sum_a A[i,a] C[k,a] to get a 3D array with indices as (i,k,a)
    AC = np.einsum('ia,ja->ija', A, C)
    
    # Calculate entire summation losing a-ith index & reshaping to desired shape
    return np.dot(AC.reshape(m*p,d),B).reshape(m,p,n).transpose(0,2,1)
    
def tensor_prod_v3(A,B,C):
    # Shape parameters
    m,d = A.shape
    n = B.shape[1]
    p = C.shape[0]
    
    # Calculate \sum_a B[a,j] C[k,a] to get a 3D array with indices as (a,j,k)
    BC = np.einsum('ai,ja->aij', B, C)
    
    # Calculate entire summation losing a-ith index & reshaping to desired shape
    return np.dot(A,BC.reshape(d,n*p)).reshape(m,n,p)

根据输入数组的形状，不同的方法会产生不同的加速，但我们希望所有方法都比 all-einsum 方法更好。下一节列出了性能数据。

运行时测试

这可能是最重要的部分，因为我们尝试使用所提出方法的三种变体来研究加速数字 all-einsum 最初在问题中提出的方法。

数据集#1（等形数组）：

In [494]: L1 = 200
     ...: L2 = 200
     ...: L3 = 200
     ...: al = 200
     ...: 
     ...: A = np.random.rand(L1,al)
     ...: B = np.random.rand(al,L2)
     ...: C = np.random.rand(L3,al)
     ...: 

In [495]: %timeit tensor_prod_v1(A,B,C)
     ...: %timeit tensor_prod_v2(A,B,C)
     ...: %timeit tensor_prod_v3(A,B,C)
     ...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
     ...: 
1 loops, best of 3: 470 ms per loop
1 loops, best of 3: 391 ms per loop
1 loops, best of 3: 446 ms per loop
1 loops, best of 3: 3.59 s per loop

数据集 #2（更大的 A）：

In [497]: L1 = 1000
     ...: L2 = 100
     ...: L3 = 100
     ...: al = 100
     ...: 
     ...: A = np.random.rand(L1,al)
     ...: B = np.random.rand(al,L2)
     ...: C = np.random.rand(L3,al)
     ...: 

In [498]: %timeit tensor_prod_v1(A,B,C)
     ...: %timeit tensor_prod_v2(A,B,C)
     ...: %timeit tensor_prod_v3(A,B,C)
     ...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
     ...: 
1 loops, best of 3: 442 ms per loop
1 loops, best of 3: 355 ms per loop
1 loops, best of 3: 303 ms per loop
1 loops, best of 3: 2.42 s per loop

数据集 #3（更大的 B）：

In [500]: L1 = 100
     ...: L2 = 1000
     ...: L3 = 100
     ...: al = 100
     ...: 
     ...: A = np.random.rand(L1,al)
     ...: B = np.random.rand(al,L2)
     ...: C = np.random.rand(L3,al)
     ...: 

In [501]: %timeit tensor_prod_v1(A,B,C)
     ...: %timeit tensor_prod_v2(A,B,C)
     ...: %timeit tensor_prod_v3(A,B,C)
     ...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
     ...: 
1 loops, best of 3: 474 ms per loop
1 loops, best of 3: 247 ms per loop
1 loops, best of 3: 439 ms per loop
1 loops, best of 3: 2.26 s per loop

数据集 #4（大 C）：

In [503]: L1 = 100
     ...: L2 = 100
     ...: L3 = 1000
     ...: al = 100
     ...: 
     ...: A = np.random.rand(L1,al)
     ...: B = np.random.rand(al,L2)
     ...: C = np.random.rand(L3,al)

In [504]: %timeit tensor_prod_v1(A,B,C)
     ...: %timeit tensor_prod_v2(A,B,C)
     ...: %timeit tensor_prod_v3(A,B,C)
     ...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
     ...: 
1 loops, best of 3: 250 ms per loop
1 loops, best of 3: 358 ms per loop
1 loops, best of 3: 362 ms per loop
1 loops, best of 3: 2.46 s per loop

数据集 #5（更大的第 a 维长度）：

In [506]: L1 = 100
     ...: L2 = 100
     ...: L3 = 100
     ...: al = 1000
     ...: 
     ...: A = np.random.rand(L1,al)
     ...: B = np.random.rand(al,L2)
     ...: C = np.random.rand(L3,al)
     ...: 

In [507]: %timeit tensor_prod_v1(A,B,C)
     ...: %timeit tensor_prod_v2(A,B,C)
     ...: %timeit tensor_prod_v3(A,B,C)
     ...: %timeit np.einsum('ia,aj,ka->ijk', A, B, C)
     ...: 
1 loops, best of 3: 373 ms per loop
1 loops, best of 3: 269 ms per loop
1 loops, best of 3: 299 ms per loop
1 loops, best of 3: 2.38 s per loop

结论： 我们看到 8x-10x 的加速比 [=38] =]问题中列出的方法。

Answer 3

我知道这现在有点老了，但这个话题经常出现。在 Matlab 中很难击败 tprod，这是一个由 Jason Farquhar 编写的 MEX 文件，可在此处获得

https://www.mathworks.com/matlabcentral/fileexchange/16275-tprod-arbitary-tensor-products-between-n-d-arrays

tprod 的工作原理与 einsum 非常相似，但它仅限于二元运算（2 个张量）。这可能不是真正的限制，因为我怀疑 einsum 只是执行一系列二进制操作。这些操作的顺序有很大的不同，我的理解是 einsum 只是按照数组传递的顺序执行它们，并且不允许多个中间产品。

tprod 也仅限于密集（完整）数组。 Kolda 的张量工具箱（在之前的 post 中提到过）确实支持稀疏张量，但其功能比 tprod 更受限制（它不允许在输出中重复索引）。我正在努力填补这些空白，但如果 Mathworks 做到了岂不是很好？

Matrix/Tensor 三重产品？

Matrix/Tensor Triple Product?

matlab

numpy

matrix

matrix-multiplication

blas

简介和解决方案代码

运行时测试