einsum方法详解（爱因斯坦求和）

einsum是pytorch、numpy中一个十分优雅的方法，如果利用得当，可完全代替所有其他的矩阵计算方法，不过这需要一定的学习成本。本文旨在详细解读einsum方法的原理，并给出一些基本示例。

一、爱因斯坦求和

爱因斯坦求和是一种对求和公式简洁高效的记法，其原则是当变量下标重复出现时，即可省略繁琐的求和符号。

比如求和公式：
$$\begin{array}{r}\sum _{i=1}^n{a}_i{b}_i=a_1{b}_1+a_2{b}_2+...+a_n{b}_n\end{array}$$
其中变量$a$与变量$b$的下标是相同的（即重复出现），则可将其记为：
$$\begin{array}{r}{a}_i{b}_i=\sum _{i=1}^n{a}_i{b}_i\end{array}$$

二、einsum方法原理

einsum方法正是利用了爱因斯坦求和简介高效的表示方法，从而可以驾驭任何复杂的矩阵计算操作。基本的框架如下：

C = einsum('ij,jk->ik', A, B)

上述操作表示矩阵A与矩阵B的点积。输入的参数分为两部分，前面表示计算操作的字符串，后面是以逗号隔开的操作对象（数量需与前面对应）。其中在计算操作表示中，"->“左边是以逗号隔开的下标索引，重复出现的索引即是需要爱因斯坦求和的；”->"右边的是最后输出的结果形式。
以上式为例，其计算公式为：$C_{ik}={\sum }_j{A}_{ij}{B}_{jk}$，其等价于矩阵A与B的点积。
这里有几条原则需要注意，之后也会和结合示例进行详解：

1. "->"左边的是对应维度，以逗号隔开
2. "->"右边的是最终output的形式
3. 如果符号"->"被省略则代表输出为整体求和
4. "…"表示省略之前或之后的所有维度
5. einsum中涉及到的计算操作有很多，包括但不限于点积、对应元素相乘、求和、转置等

三、一些示例

einsum方法在numpy和pytorch中均有内置，这里以pytorch为例，首先定义一些需要用到的变量：

import torch
from torch import einsum
a = torch.ones(3,4)
b = torch.ones(4,5)
c = torch.ones(6,7,8)
d = torch.ones(3,4)
x, y = torch.randn(5), torch.randn(5)

1. 计算矩阵所有元素之和

einsum('i,j', a)   # 等价于einsum('i,j->', a)
einsum('i,j,k', c)

2. 计算矩阵的迹

einsum('ii', a)

3. 获取矩阵对角线元素组成的向量

einsum('ii->i', a)

4. 向量相乘得到矩阵

einsum('i,j->ij', x, y)

5. 矩阵点积

einsum('ij,jk->ik', a, b)

6. 矩阵对应元素相乘

einsum('ij,ij->ij', a, d)

7. 矩阵的转置

einsum('ijk->ikj', c)
einsum('...jk->...kj', c)  # 两种形式等价

8. 双线性运算

A = torch.randn(3,5,4)
l = torch.randn(2,5)
r = torch.randn(2,4)
torch.einsum('bn,anm,bm->ba', l, A, r)

9. 最后来一个复杂的

a = torch.randn(3,4,5)
b = torch.randn(6,5)
c = torch.randn(6,3)
target = einsum('fti,di,df->dt', a,b,c)

# 把上面的求和过程写成循环的形式方便理解
# 对f和i进行求和
dt = torch.zeros(6,4)
for d in range(6):
	for t in range(4):
		tmp = 0
		for f in range(3):
			for i in range(5):
				tmp += a[f,t,i].item() * b[d,i].item() * c[d,f].item()
		dt[d,t] = tmp

print(target)
print('\n', dt)

# 结果：
tensor([[ 0.5712, -0.9473,  0.1972,  0.6474],
        [ 0.6273, -1.3820,  2.0779, -0.5729],
        [ 4.0083, -2.9979,  0.9493,  0.8467],
        [15.9985, -7.1897,  1.1957,  2.7394],
        [ 0.4298, -0.6509,  0.5128,  1.7027],
        [-4.6843, -4.7948,  1.1389, -6.9389]])
        
tensor([[ 0.5712, -0.9473,  0.1972,  0.6474],
        [ 0.6273, -1.3820,  2.0779, -0.5729],
        [ 4.0083, -2.9979,  0.9493,  0.8467],
        [15.9985, -7.1897,  1.1957,  2.7394],
        [ 0.4298, -0.6509,  0.5128,  1.7027],
        [-4.6843, -4.7948,  1.1389, -6.9389]])