【Pytorch】科研代码技巧—Einsum爱因斯坦求和约定

不知大家在看论文代码的时候是否会常常看见 torch.einsum()，这玩意儿看起来是真的抽象，但是深入了解后发现它原来这么好用

一、Introeduction 

einsum真名叫做爱因斯坦求和约定，用于简洁的表示转置、内积、外积、各种求和

先看看以下这个例子，有这样的一段代码

R=torch.einsum('ik,jk->ij',A,B)

 在Einsum中，箭头从左边到右边消失了什么参数，那公式前就加一个带什么参数的求和符。本案例中消失了k，因此我们需要在加上对带k的求和符，转化为数学公式如下

对数学敏感的小伙伴可能已经知道这个公式代表什么意思了。为了更好展示出来，便将其可视化

画完图之后我们可以直观的知道这就是将两个矩阵每行向量进行求内积

此外，它的底层代码就是套了很多层的for循环，如果我们不用Einsum来实现以上的功能估计要写半天了

二、Skill

是不是感觉上面的公式很简单神奇，是的Einsum诞生的初衷就是为了简化矩阵的运算，因此博主记录了以下几个常用的矩阵运算用Einsum来实现。假设有以下四个矩阵

 2.1 求某行、列、维度之和

# 行之和
R=torch.einsum('ij->i',A)

# 列之和
R=torch.einsum('ij->j',A)

# 某维度之和
R=torch.einsum('ijklmn->n',D)

2.2 所有元素之和

# 所有元素之和
R=torch.einsum('ijklmn->',D)

2.3 转置

# 转置
R=torch.einsum('ij->ji',A)

2.4 内积

# 内积
R=torch.einsum('ij,jk->ik',A,B)

2.5 外积

# 外积
R=torch.einsum('ij,ik->jk',A,C)

2.6 灵活相乘 

# 多维相乘
R=torch.einsum('ij,jk,lj->jk',A,B,C)