einsum函数

einsum函数是什么

einsum就是爱因斯坦求和约定，其实质上是在求和过程中将求和符号省略, 在张量分析中求和总是以上下重复指标的形式出现，因此这些指标是傀儡指标，即使去掉也不影响, i.e.,

$$C_i^{\cdot j}=A_i^{\cdot k}{B}_k^{\cdot j}$$

这其实是一个矩阵乘法, 如果你愿意至少可以写成有三种实现方式.

• For Loop

I, K = A.size()
K, J = B.size()
for i in range(I):
    for j in range(J):
        for k in range(K):
            C[i,j] += A[i,k]*B[k,j]

• 矩阵乘法

C = torch.matmul(A,B)

• einsum

C = torch.einsum("ik,kj->ij", A, B) 

具体看一下用einsum是如何工作的

对于 C = torch.einsum(“ij,jk->ik”, A, B)，C是运行的结果，A，B是进行计算的参数，“ij,jk->ik”给出了维度信息. 这一计算过程就是如何利用A[i,k]和B[j,k]得到C[i,j].

我们要做的就是补上$\sum$符号，使得张量的einsum运算可以被公式化，而后让我们理解。

具体的补充操作要看左右两边指标的差异，i.e.,

{ i , k } ∪ { k , j } − { i , j } = { k } \left\{ i,k\right\}\cup\left\{ k,j\right\}-\left\{ i,j\right\} = \left\{ k\right\} {i,k}∪{k,j}−{i,j}={k}

集合的差集为 { k } \left\{ k\right\} {k}那就补上${\sum }_k$, i.e.,

$$C_{i,j}=\sum _k{A}_{i,k}{B}_{k,j}$$

Remak: 因为这里是笛卡尔坐标系，上下标是一样的，但是指标的前后顺序不能变.

einsum除了计算矩阵乘法以外还有一些别的用法

• 转置

B = torch.einsum('ij->ji', A)

• 求行和

B = torch.einsum('ij->j', A)

• 求列和

B = torch.einsum('ij->i', A)

• 求和至标量

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

Remark: 因为标量没有指标所以就是没有.

• 矩阵向量乘法

C = torch.einsum('ik,k->i', A, B)

• 向量内积

C= torch.einsum('i,i->', A, B)

• 向量外积

C = torch.einsum('i, j->ij', A, B)

• Hardamard 积

C = torch.einsum('ij,ij->ij', A, B)

以上还都只是在矩阵和向量范围内，其实einsum的真正的用武之地是在张量的缩并上。

• 双线性变换

$$D_{ij}=\sum _k\sum _l{A}_{ik}{B}_{jkl}{C}_{il}=A_{ik}{B}_{jkl}{C}_{il}$$

D = torch.einsum('ik,jkl,il->ij', A, B. C)

• 张量缩并

$$C_{klmn}=\sum _{ij}{A}_{ijkl}{B}_{ijmn}$$

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)

Reamrk: 在我看来真正的难点在于思维的转换，我们之前学习线性代数或统计机器学习往往都会将公式写成矩阵乘法的简洁形式，但这是写法只在二维以下有效。真正在高维的张量的运算中都是对张量元素的操作。只要转变了思维，不拘泥于整体形式，转而关心张量各个分量，就会发现einsum的写法是非常整洁和高效。

参考文献

1. https://blog.csdn.net/zhaohongfei_358/article/details/125273126?spm=1001.2014.3001.5502
2. https://rogerspy.github.io/2021/09/12/einsum-mhsa/