爱因斯坦求和—— einsum函数用法（pytorch）

爱因斯坦求和—— einsum函数用法（pytorch）

---

---

前言

这个函数也是偶然间遇到的，觉得还挺好使的。

---

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

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

->符号就相当于等号，->左边为输入，右边为输出。
还是直接看程序吧。

导入包

# 导入包
import torch
from torch import einsum

求和

a = torch.Tensor(range(2*3)).view(2, 3)
b = torch.einsum("ij -> i", a) #按行求和
c = torch.einsum("ij -> j", a) #按列求和
d = torch.einsum("ij -> ", a) #全部求和
a,b,c,d

# 求和
# a为3维 b维2维   下面就是把a按照第维度0对应点求和得出b
# 等价于 b = a.sum(dim=0)
a = torch.Tensor(range(2*3*4)).view(2, 3, 4)
b = torch.einsum("ijk->jk", a)
a.shape,b.shape,a,b

求转置矩阵

a = torch.Tensor(range(2*3)).view(2, 3)
b = torch.einsum("ij->ji", a)
a,b

矩阵相乘

# 矩阵相乘
a = torch.Tensor(range(2*3)).view(2, 3)
b = torch.Tensor(range(3*4)).view(3, 4)
c = torch.einsum("ij,jk->ik", a, b) 
d = a@b
a,b,c,d

求矩阵的迹

# 求矩阵的迹
a = torch.randn(4, 4)
b = torch.einsum('ii', a)
a,b

提取矩阵对角线元素

# 提取矩阵对角线元素为向量
a = torch.randn(4, 4)
b = torch.einsum('ii->i',a )
a,b

内积 外积

# 张量内积 外积
a = torch.tensor([1,2])
b = torch.tensor([1,3,5])
c = torch.tensor([3,4])

## 内积
d = torch.einsum("i, j ->", a, c)
## 外积
e = torch.einsum("i, j -> ij", a, b)
d,e

求批矩阵乘法

# 求批矩阵乘法
a = torch.randn(3,2,5)
b = torch.randn(3,5,4)
c = torch.einsum('bij,bjk->bik', [As, Bs])
a.shape,b.shape,c.shape

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)
a.shape,b.shape,c.shape

---

总结

未完待续，，，，