matlab - 矩阵/张量三重积？

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

简介及解决方案代码

np.einsum , 确实很难被打败，但是在极少数情况下，如果您可以将矩阵乘法引入计算中，您仍然可以打败它。试了几次好像可以带进来matrix-multiplication with np.dot用 np.einsum('ia,aj,ka->ijk', A, B, C) 超越性能。

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

• A:[i,a] 和 B:[a,j] 的求和是用 np.einsum 得到的我们是一个 3D 数组:[i,j,a]。
• 然后将此 3D 数组重新整形为 2D 数组:[i*j,a] 并将第三个数组 C[k,a] 转置为 [a,k]，目的是在这两者之间执行矩阵乘法，给我们[i*j,k] 作为矩阵产品，因为我们在那里丢失了索引 [a]。
• 产品被 reshape 为 3D 数组:[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 的加速比 all-einsum 所提议方法的变体 问题中列出的方法。