文章目录
- 本文内容
- Einsum函数简介
- 如何看懂一个einsum式子
- 如何看懂一个einsum式子(实战)
- einsum特殊写法补充
- 如何写出einsum表达式
本文内容
可能你在某个地方听说了einsum,然后不会写,或者看不懂。这篇文章将会一步一步教会你如何使用(通法哦,只要学会方法就全会了)。
Einsum函数简介
ein 就是爱因斯坦的ein,sum就是求和。einsum就是爱因斯坦求和约定,其实作用就是把求和符号省略,就这么简单。举个例子:
我们现在有一个矩阵
我们想对A的“行”进行求和得到矩阵B(向量B),用公式表示,则为:
对于这个求和符号,爱因斯坦说看着有点多余,要不就省略了吧,然后式子就变成了:
用einsum表示呢,则为: torch.einsum("ij->i", A)
。->
符号就相当于等号,->
左边的ij
就相当于,
->
右边的i
就相当于。
einsum
接收的第一个参数为einsum表达式,后面的参数为等号右边的矩阵。
不只是pytorch里有,numpy,tensonflow这些里面都有einsum。
这里的是指代A的下标,也可以换成其他字母
到这里,如果悟性好的同学应该就已经彻底懂了。但应该还有很多同学和我一样处于懵逼状态,所以接下来我会讲解如何看懂一个einsum公式和如何写出einsum表达式。
如何看懂一个einsum式子
当我们拿到一个einsum表达式后,第一步是要写出它的数学表达式。例如,我们有如下一个einsum表达式:
A = torch.Tensor(range(2*3*4)).view(2, 3, 4)
C = torch.einsum("ijk->jk", A)
则,该式子的数学表达式为:
第二步,补充符号,那如何补,补几个,
下面放什么呢?这里就要看左右两边下标的差异了,要补的
符号就是右边的下标减左边的下标。在这个例子中,右边有
,而左边是
,差了一个
,所以补
。最终为:
第三步,用笔纸画出(或脑补出)这个等式到底干了些啥,对于该等式,可以画为:
这样就可以很容易看出来,它是将行都给加一起了,等价于
C = A.sum(dim=0)
第四步,尝试用for循环复现,其实einsum还是很好复现的,就按照公式写for循环就行了,求和的部分用+=
i, j, k = A.shape[0], A.shape[1], A.shape[2] # 得到 i, j, k
C_ = torch.zeros(j, k) # 初始化 C_ , 用来保存结果
for i_ in range(i): # 遍历 i
for k_ in range(k): # 遍历 j
for j_ in range(j): # 遍历 k
C_[j_][k_] += A[i_][j_][k_] # 求和
C, C_
(tensor([[12., 14., 16., 18.],
[20., 22., 24., 26.],
[28., 30., 32., 34.]]),
tensor([[12., 14., 16., 18.],
[20., 22., 24., 26.],
[28., 30., 32., 34.]]))
可以看到,我们的for循环结果和einsum的结果一致。
到这里,如何看懂einsum就结束了,按照上面四步走,多加练习即可。
如何看懂一个einsum式子(实战)
我也练几个。先来一个简单的。
A = torch.Tensor(range(2*3)).view(2, 3)
B = torch.einsum("ij->ji", A)
第一步,写出数学表达式:
第二步,添加符号,这里左边是
,右边是
,不多不少,正正好,所以不需要(也不能)增添
符号。
第三步,画出矩阵的变换过程:
哦,这不就是求转置矩阵嘛。
第四步,使用for循环复现:
i, j = A.shape[0], A.shape[1] # 得到 i, j
B_ = torch.zeros(j, i) # 初始化 B_ , 用来保存结果
for i_ in range(i): # 遍历 i
for j_ in range(j): # 遍历 j
B_[j_][i_] = A[i_][j_] # 因为不需要求和,所以这里用=,而不是+=“”
B, B_
(tensor([[0., 3.],
[1., 4.],
[2., 5.]]),
tensor([[0., 3.],
[1., 4.],
[2., 5.]]))
接下来来个难的。
A = torch.Tensor(range(2*3*4*5)).view(2, 3, 4, 5)
B = torch.Tensor(range(2*3*7*8)).view(2, 3, 7, 8)
C = torch.einsum("ijkl,ijmn->klmn", A, B)
如果等式右边有多个矩阵,则用逗号分割。
第一步,写出数学表达式:
第二步,补充求和符号,右边有,左边有
,左边少了
,所以补两个求和符号,即
。最终为:
注意这里可不是矩阵相乘,而是两个数字相乘,因为
和
都是数字
第三步,画出矩阵变换过程。四维太难画了,脑补吧。
第四步,使用for循环进行复现。
i,j,k,l,m,n = A.shape[0],A.shape[1],A.shape[2],A.shape[3],B.shape[2],B.shape[3]
C_ = torch.zeros(k,l,m,n)
for i_ in range(i):
for j_ in range(j):
for k_ in range(k):
for l_ in range(l):
for m_ in range(m):
for n_ in range(n):
# 由于有求和符号,所以用+=
C_[k_][l_][m_][n_] += A[i_][j_][k_][l_]*B[i_][j_][m_][n_]
C == C_
tensor([[[[True, True, True, ..., True, True, True],
...........................
[True, True, True, ..., True, True, True]]]])
einsum特殊写法补充
- 若等号左边就是一个数,那么
->
左边什么都不用写,例如:
A = torch.Tensor(range(1*2*3)).view(1, 2, 3)
b = torch.einsum("ijk->", A) # 由于b是一个数,没有下标,所以->右边什么都不用写
b
tensor(15.)
- 若下标过多,或不确定,则可以省略,例如:
A = torch.Tensor(range(1*2*3)).view(1, 2, 3)
B = torch.einsum("i...->...", A) # 省略号表示*
B.size()
torch.Size([2, 3])
目前为止,你应该可以看得懂einsum表达式了,若看不懂,大概率是因为公式的问题,确实有些求和公式很复杂,你可以慢慢拆解求和公式,看看具体表示的什么含义。
如何写出einsum表达式
要写出einsum表达式也很简单,只要将上面的步骤反过来就行了,①先画出你要做的矩阵运算;②尝试用for循环实现;③写出数学表达式;④写出einsum表达式,并验证。
接下来,我们用矩阵相乘公式来进行演示。第一步,我们要画出矩阵相乘的操作过程,如下:
第二步,尝试使用for循环实现:
A = torch.Tensor(range(2*3)).view(2, 3)
B = torch.Tensor(range(3*4)).view(3, 4)
C = torch.zeros(i, k)
i, j, k = 2, 3, 4
for i_ in range(i):
for j_ in range(j):
for k_ in range(k):
C[i_][k_] += A[i_][j_]*B[j_][k_]
第三步,写出数学表达式:
第3.2步,补充求和符号,左边是,右边是
,少了
,补
:
第四步,写出einsum表达式并验证:
D = torch.einsum("ij,jk->ik", A, B)
E = A@B
C, D, E
(tensor([[20., 23., 26., 29.],
[56., 68., 80., 92.]]),
tensor([[20., 23., 26., 29.],
[56., 68., 80., 92.]]),
tensor([[20., 23., 26., 29.],
[56., 68., 80., 92.]]))
参考资料:
einsum is all you need: https://www.youtube.com/watch?v=pkVwUVEHmfI