einsum pytorch和numpy

np.einsum和torch.einsum均使用一致的签名einsum(equation, operands)

   其中equation是表示爱因斯坦求和约定字符串，而operands则是张量序列（numpy和tensorflow中是变长参数列表，而在pytorch中是列表）cj=∑i∑kAikBkj写成equation字符串就是ik,kj -> j。注意这里(i, j, k)的命名是任意的，但需要一致。
   PyTorch和TensorFlow像numpy支持einsum的好处之一是einsum可以用于神经网络架构的任意计算图，并且可以反向传播。典型的einsum调用格式如下：
   ![在这里插入图片描述](https://img-blog.csdnimg.cn/684ed23264e342f2834399077b21b249.jpg#pic_center)

上式中◻是占位符，表示张量维度。上面的例子中，arg1和arg3是矩阵，arg2是二阶张量，这一einsum运算的结果（result）是矩阵。注意einsum处理的是可变数量的输入。在上面的例子中，einsum指定了三个参数之上的操作，但它同样可以用在牵涉一个参数、两个参数、三个以上参数的操作上。学习einsum的最佳途径是通过学习一些例子，所以下面我们将展示一下，在许多深度学习模型中常用的库函数，用einsum该如何表达（以PyTorch为例）。
2.1 矩阵转置

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

2.2 求和

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

2.3 列求和

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

2.4 行求和

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

2.5 矩阵-向量相乘

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

2.6 矩阵-矩阵相乘

a = torch.arange(6).reshape(2, 3)
b = torch.arange(15).reshape(3, 5)
torch.einsum('ik,kj->ij', [a, b])
tensor([[  25.,   28.,   31.,   34.,   37.],
        [  70.,   82.,   94.,  106.,  118.]])

2.7 点积

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

矩阵：

a = torch.arange(6).reshape(2, 3)
b = torch.arange(6,12).reshape(2, 3)
torch.einsum('ij,ij->', [a, b])
tensor(145.)

2.8 哈达玛积

a = torch.arange(6).reshape(2, 3)
b = torch.arange(6,12).reshape(2, 3)
torch.einsum('ij,ij->ij', [a, b])
tensor([[  0.,   7.,  16.],
        [ 27.,  40.,  55.]])

2.9 外积

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

2.10 batch矩阵相乘

a = torch.randn(3,2,5)
b = torch.randn(3,5,3)
torch.einsum('ijk,ikl->ijl', [a, b])
tensor([[[ 1.0886,  0.0214,  1.0690],
         [ 2.0626,  3.2655, -0.1465]],

        [[-6.9294,  0.7499,  1.2976],
         [ 4.2226, -4.5774, -4.8947]],

        [[-2.4289, -0.7804,  5.1385],
         [ 0.8003,  2.9425,  1.7338]]])

2.11 张量缩约

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])

2.12 双线性变换
如前所述，einsum可用于超过两个张量的计算。这里举一个这方面的例子，双线性变换。

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])
tensor([[ 3.8471,  4.7059, -3.0674, -3.2075, -5.2435],
        [-3.5961, -5.2622, -4.1195,  5.5899,  0.4632]])

einsum满足你一切需要：深度学习中的爱因斯坦求和约定