einsum：爱因斯坦求和约定

深藏阁楼爱情的钟 2022-11-28 00:47 335阅读 0赞

在Tensorflow、Numpy和PyTorch中都提供了使用einsum的api，einsum是一种能够简洁表示点积、外积、转置、矩阵-向量乘法、矩阵-矩阵乘法等运算的领域特定语言。在Tensorflow等计算框架中使用einsum，操作矩阵运算时可以免于记忆和使用特定的函数，并且使得代码简洁，高效。  
如对矩阵 A ∈ R I × K A∈R^\{I\\times K\} A∈RI×K和矩阵 B ∈ R K × J B∈R^\{K\\times J\} B∈RK×J做矩阵乘，然后对列求和，最终得到向量 c ∈ R J c∈R^J c∈RJ，即：  
![20200816175653436.png_pic_center][]  
使用爱因斯坦求和约定表示为：  
![在这里插入图片描述][20200816175725231.png_pic_center]  
在Tensorflow、Numpy和PyTorch中对应的einsum字符串为：`ik,kj->j`。在上面的字符串中，隐式地省略了重复的下标k，表示在该维度矩阵乘；另外输出中未指明下标i，表示在该维度累加。

## Numpy、PyTorch和Tensorflow中的einsum ##

einsum在Numpy中的实现为np.einsum，在PyTorch中的实现为torch.einsum，在Tensorflow中的实现为tf.einsum，均使用同样的函数签名einsum(equation,operands)，其中，equation传入爱因斯坦求和约定的字符串，而operands则是张量序列。在Numpy、Tensorflow中是变长参数列表，而在PyTorch中是列表。上述例子中，在Tensorflow中可写作：`tf.einsum('ik,kj->j',mat1,mat2)`。 其中，mat1、mat2为执行该运算的两个张量。注意：这里的(i,j,k)的命名是任意的，但在一个表达式中要一致。  
PyTorch和Tensorflow像Numpy支持einsum的好处之一就是，einsum可以用于深度网络架构的任意计算图，并且可以反向传播。在Numpy和Tensorflow中的调用格式如下：![在这里插入图片描述][20200816180404150.png_pic_center]  
其中，□是占位符，表示张量维度；arg1,arg3是矩阵，arg2是三阶张量，运算结果是矩阵。注意：einsum处理可变数量的输入。上面例子中，einsum制定了三个参数的操作，但同样可以操作一个参数、两个参数和三个参数及以上的操作。

## 典型的einsum表达式 ##

内积：又称点积、点乘，对应位置数字相乘，结果是一个标量，有见向量内积和矩阵内积等。  
向量a和向量b的内积：  
![在这里插入图片描述][20200816180732718.png_pic_center]  
内积几何意义：![在这里插入图片描述][2020081618075958.png_pic_center]  
外积：又称叉乘、叉积、向量积，行向量矩阵乘列向量，结果是二阶张量。注意到：张量的外积作为张量积的同义词。外积是一种特殊的克罗内克积。  
向量a和向量b 的外积：  
![在这里插入图片描述][20200816180904293.png_pic_center]  
外积的几何意义：  
![在这里插入图片描述][202008161809387.png_pic_center]  
其中，  
![在这里插入图片描述][20200816181000600.png_pic_center]  
由于PyTorch可以实时输出运算结果，以PyTorch使用einsum表达式为例。

矩阵转置 Bji=Aij

a=torch.arange(6).reshape(2,3)
    >>>tensor([[0, 1, 2],
               [3, 4, 5]])
    torch.einsum('ij->ji',[a])
    >>>tensor([[0, 3],
              [1, 4],
              [2, 5]])

求和  b = ∑ ∑ A i j b=\\sum\{\\sum\{A\_\{ij\}\}\} b=∑∑Aij

a=torch.arange(6).reshape(2,3)
    >>>tensor([[0, 1, 2],
               [3, 4, 5]])
    torch.einsum('ij->',[a])
    >>>tensor(15)

列求和(列维度不变，行维度消失)  b j = ∑ A i j b\_j=\\sum\{A\_\{ij\}\} bj=∑Aij

a=torch.arange(6).reshape(2,3)
    >>>tensor([[0, 1, 2],
               [3, 4, 5]])
    torch.einsum('ij->j',[a])
    >>>tensor([ 3.,  5.,  7.])

列求和(列维度不变，行维度消失)  b i = ∑ A i j b\_i=\\sum\{A\_\{ij\}\} bi=∑Aij

a=torch.arange(6).reshape(2,3)
    >>>tensor([[0, 1, 2],
               [3, 4, 5]])
    torch.einsum('ij->i', [a])
    >>>tensor([  3.,  12.])

矩阵-向量相乘  
![在这里插入图片描述][20200816181629137.png_pic_center]

a=torch.arange(6).reshape(2,3)
    >>>tensor([[0, 1, 2],
               [3, 4, 5]])
    torch.einsum('ik,k->i',[a,b])
    >>>tensor([  5.,  14.])

矩阵-矩阵乘法  
![在这里插入图片描述][20200816181711578.png_pic_center]

a=torch.arange(6).reshape(2,3)
    b=torch.arange(15).reshape(3,5)
    >>>tensor([[0, 1, 2],
               [3, 4, 5]])
               
    >>>tensor([[ 0,  1,  2,  3,  4],
               [ 5,  6,  7,  8,  9],
               [10, 11, 12, 13, 14]])
    torch.einsum('ik,kj->ij',[a,b])
    >>>tensor([[ 25,  28,  31,  34,  37],
               [ 70,  82,  94, 106, 118]])

点积  
向量  
![在这里插入图片描述][20200816181801345.png_pic_center]

a=torch.arange(3)
    b=torch.arange(3,6)
    >>>tensor([0, 1, 2])
    >>>tensor([3, 4, 5])
    torch.einsum('i,i->',[a,b])
    >>>tensor(14.)

矩阵  
![在这里插入图片描述][20200816181836236.png_pic_center]

a=torch.arange(6).reshape(2,3)
    b=torch.arange(6,12).reshape(2,3)
    >>>tensor([[0, 1, 2],
               [3, 4, 5]])
    
    >>>tensor([[ 6,  7,  8],
               [ 9, 10, 11]])
    torch.einsum('ij,ij->',[a,b])
    >>>tensor(145.)

外积  
![在这里插入图片描述][20200816181917906.png_pic_center]

a=torch.arange(3)
    b=torch.arange(3,7)
    >>>tensor([0, 1, 2])
    >>>tensor([3, 4, 5, 6])
    torch.einsum('i,j->ij',[a,b])
    >>>tensor([[  0.,   0.,   0.,   0.],
               [  3.,   4.,   5.,   6.],
               [  6.,   8.,  10.,  12.]])

batch矩阵乘  
![在这里插入图片描述][20200816181953677.png_pic_center]

a=torch.randn(3,2,5)
    b=torch.randn(3,5,3)
    >>>tensor([[[-1.4131e+00,  3.8372e-02,  1.2436e+00,  5.4757e-01,  2.9478e-01],
                [ 1.3314e+00,  4.4003e-01,  2.3410e-01, -5.3948e-01, -9.9714e-01]],
    
               [[-4.6552e-01,  5.4318e-01,  2.1284e+00,  9.5029e-01, -8.2193e-01],
                [ 7.0617e-01,  9.8252e-01, -1.4406e+00,  1.0071e+00,  5.9477e-01]],
    
              [[-1.0482e+00,  4.7110e-02,  1.0014e+00, -6.0593e-01, -3.2076e-01],
               [ 6.6210e-01,  3.7603e-01,  1.0198e+00,  4.6591e-01, -7.0637e-04]]])
    
    >>>tensor([[[-2.1797e-01,  3.1329e-04,  4.3139e-01],
                [-1.0621e+00, -6.0904e-01, -4.6225e-01],
                [ 8.5050e-01, -5.8867e-01,  4.8824e-01],
                [ 2.8561e-01,  2.6806e-01,  2.0534e+00],
                [-5.5719e-01, -3.3391e-01,  8.4069e-03]],
    
               [[ 5.2877e-01,  1.4361e+00, -6.4232e-01],
                [ 1.0813e+00,  8.5241e-01, -1.1759e+00],
                [ 4.9389e-01, -1.7523e-01, -9.5224e-01],
                [-1.3484e+00, -5.4685e-01,  8.5539e-01],
                [ 3.7036e-01,  3.4368e-01, -4.9617e-01]],
    
               [[-2.1564e+00,  3.0861e-01,  3.4261e-01],
                [-2.3679e+00, -2.5035e-01,  1.8104e-02],
                [ 1.1075e+00,  7.2465e-01, -2.0981e-01],
                [-6.5387e-01, -1.3914e-01,  1.5205e+00],
                [-1.6561e+00, -3.5294e-01,  1.9589e+00]]])
    torch.einsum('ijk,ikl->ijl',[a,b])
    >>>tensor([[[ 1.3170, -0.7075,  1.1067],
                [-0.1569, -0.2170, -0.6309]],
    
               [[-0.1935, -1.3806, -1.1458],
                [-0.4135,  1.7577,  0.3293]],
    
               [[ 4.1854,  0.5879, -2.1180],
                [-1.4922,  0.7846,  0.7267]]])

张量缩约  
batch矩阵相乘是张量缩约的一个特例，比如有两个张量，一个n阶张量 A ∈ R I 1 × l 2 × . . . × I n  A∈R^\{I\_1\\times l\_2 \\times ... \\times I\_n\} A∈RI1×l2×...×In，一个m阶张量 B ∈ R J 1 × J 2 × . . . × J m  B∈R^\{J\_1 \\times J\_2 \\times...\\times J\_m\} B∈RJ1×J2×...×Jm。取n=4，m=5，假定维度 I 2 = J 3 I\_2=J\_3 I2=J3且 I 3 = J 5  I\_3=J\_5 I3=J5，将这两个张量在这两个维度上(A张量的第2、3维度，B张量的第3、5维度)相乘，获得新张量 C ∈ R I 1 × I 4 × J 1 × J 2 × J 4 C∈R^\{I\_1 \\times I\_4 \\times J\_1 \\times J\_2 \\times J\_4\} C∈RI1×I4×J1×J2×J4，如下所示：  
![在这里插入图片描述][20200816182050728.png_pic_center]

a=torch.randn(2,3,5,7)
    b=torch.randn(11,13,3,17,5)
    
    torch.einsum('pqrs,tuqvr->pstuv', [a, b]).shape
    >>>torch.Size([2, 7, 11, 13, 17])

多张量计算  
如前所述，einsum可用于超过两个张量的计算，以双线性变换为例：  
![在这里插入图片描述][20200816182557974.png_pic_center]

a=torch.randn(2,3)
    b=torch.randn(5,3,7)
    c=torch.randn(2,7)
    
    torch.einsum('ik,jkl,il->ij',[a,b,c]).shape
    >>>torch.Size([2,5])

[20200816175653436.png_pic_center]: /images/20221124/a80519e9be7d4ee7ac8332a14e221416.png
[20200816175725231.png_pic_center]: /images/20221124/b7c4718b3a1c49c988459c25408c7b37.png
[20200816180404150.png_pic_center]: /images/20221124/41b59d7dfc134c1e932365fd55fa55dd.png
[20200816180732718.png_pic_center]: /images/20221124/49120a241fef4bbe9e31e56ce5a9c970.png
[2020081618075958.png_pic_center]: /images/20221124/6e09d5057f9d4f419ef0da5eed46840d.png
[20200816180904293.png_pic_center]: /images/20221124/ece220a5f74f46e2ad87afc04bfc2af6.png
[202008161809387.png_pic_center]: /images/20221124/e8a65c5497bd41dfb17ee936c2e506f6.png
[20200816181000600.png_pic_center]: /images/20221124/b7ff1868f8994638afcd8034e720ca9b.png
[20200816181629137.png_pic_center]: /images/20221124/a193a3afa7954e9399ad6293fa9b48f9.png
[20200816181711578.png_pic_center]: /images/20221124/abffd523be834c62ae4e2075cf62d772.png
[20200816181801345.png_pic_center]: /images/20221124/c7fe7929aa784831bb52e7d3f1edfeca.png
[20200816181836236.png_pic_center]: /images/20221124/70cc52b2b6ad4dca8cde66ddf5726ff3.png
[20200816181917906.png_pic_center]: /images/20221124/f1233874326049ddb3b352dc9c25acbb.png
[20200816181953677.png_pic_center]: /images/20221124/2e1e55363d4e4660a7d1132986948a3f.png
[20200816182050728.png_pic_center]: /images/20221124/003ad98d451e49fb927223c76b59d602.png
[20200816182557974.png_pic_center]: /images/20221124/f817e32b985f430ca5bccd61ab2b9afd.png