PyTorch——通俗理解torch.einsum

参考链接

1. https://www.cnblogs.com/mengnan/p/10319701.html

从einsum表达式恢复为张量计算

torch.einsum一般这样用：

$result=torch.einsum\left({}^{\prime }\text{xx,xx,xx}\to {\text{xx}}^{\prime },arg1,arg2,arg3\right)$

问题在于如何把里面的einsum表达式${}^{\prime }\text{xx,xx,xx}\to {\text{xx}}^{\prime },arg1,arg2,arg3$，还原为直观的张量计算形式。

里面的$arg$为输入的张量，$\to$左边的$\text{xx,xx,xx}$为输入张量的下标，如$\text{ik,jkl,il}$，有多少个输入就有多少个下标，$\to$右边的$\text{xx}$为输出张量的下标，有时候也空着。恢复步骤如下：

1. 首先写成输出的模板，用$R$表示输出张量。

$R_{\text{xx}}=arg{1}_{\text{xx}}\cdot arg{2}_{\text{xx}}\cdot arg{3}_{\text{xx}}$

2. 将输入张量下标（即$\to$左边的$\text{xx,xx,xx}$）中没有在输出张量下标（即$\to$右边的$\text{xx}$）出现的字母，用求和符号加在等式右边。

$R_{\text{xx}}=\sum \sum arg{1}_{\text{xx}}\cdot arg{2}_{\text{xx}}\cdot arg{3}_{\text{xx}}$

实例

1. torch.einsum(‘ijk,ikl->ijl’,[a,b])

第一步：

$R_{ijl}=A_{ijk}{B}_{ikl}$

第二步：
$k$没有在输出张量下标中出现，所以：

$R_{ijl}={\sum }_k{A}_{ijk}{B}_{ikl}$

2. torch.einsum(‘ij->i’, [a])

第一步：

$R_i=A_{ij}$

第二步：
$j$没有在输出张量下标中出现，所以：

$R_i={\sum }_j{A}_{ij}$

3. torch.einsum(‘ij->’,[a])

第一步：

$R=A_{ij}$

第二步：

$R={\sum }_i{\sum }_j{A}_{ij}$