Pytorch中, torch.einsum详解。

爱因斯坦简记法：是一种由爱因斯坦提出的，对向量、矩阵、张量的求和运算的求和简记法。

在该简记法当中，省略掉的部分是：1）求和符号与2）求和号的下标 

省略规则为：默认成对出现的下标（如下例1中的i和例2中的k）为求和下标。

我们举例说明：

 

1）用简化表示内积

其中o为输出。

2) 用简化表示矩阵乘法

其中为输出矩阵的第ij个元素。

这样的求和简记法，能够以一种统一的方式表示各种各样的张量运算（内积、外积、转置、点乘、矩阵的迹、其他自定义运算），为不同运算的实现提供了一个统一模型。

因此，numpy和pytorch对它进行了实现。两者的语义一致，本文各举一例。

3）Torch矩阵乘法。

print(a_tensor)
 
tensor([[11, 12, 13, 14],
        [21, 22, 23, 24],
        [31, 32, 33, 34],
        [41, 42, 43, 44]])

print(b_tensor)

tensor([[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3],
        [4, 4, 4, 4]])

# 'ik, kj -> ij'语义解释如下：
# 输入a_tensor: 2维数组，下标为ik,
# 输入b_tensor: 2维数组，下标为kj,
# 输出output：2维数组，下标为ij。
# 隐含语义：输入a,b下标中相同的k，是求和的下标，对应上面的例子2的公式
output = torch.einsum('ik, kj -> ij', a_tensor, b_tensor)

print(output)

tensor([[130, 130, 130, 130],
        [230, 230, 230, 230],
        [330, 330, 330, 330],
        [430, 430, 430, 430]])

2）高阶张量运算的一个例子：

a = np.arange(60.).reshape(3,4,5)
b = np.arange(24.).reshape(4,3,2)

# 语义解析：
# 输入a：3阶张量，下标为ijk
# 输入b: 3阶张量，下标为jil
# 输出o: 2阶张量，下标为k和l
# 隐含语义：对i,j进行求和，公式附于代码之后：
o = np.einsum('ijk,jil->kl', a, b)
print(o)

array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

# 验证：
print(np.sum(a[:,:,0]*b[:,:,0].T))

4400.0

print(np.sum(a[:,:,1]*b[:,:,0].T))

4532.0

其公式为：

其语义为：O第k，l个元素：是矩阵A[:,:,k]和矩阵B[:,:,l]转置，对应元素相乘再求和。

 

其他例子请参考https://www.cnblogs.com/mengnan/p/10319701.html