【Pytorch】einsum爱因斯坦简记法举例

95 阅读 0 评论 63 点赞

我是靠谱客的博主高高枫叶，最近开发中收集的这篇文章主要介绍【Pytorch】einsum爱因斯坦简记法举例，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

本文为 Pytorch 学习笔记，讲解爱因斯坦简记法。欢迎在评论区与我讨论????

爱因斯坦简记法

是爱因斯坦提出的对向量、矩阵、张量求和运算的简记法。

被省略的部分是求和符号和求和号的下标。

省略规则：默认成对出现的下标为求和下标。

举例

用 $x_i,y_i$ 简化表示内积 $< x, y >$ ：
$x_iy_i:=sum_ix_iy_i=o$
其中 $o$ 为输出。
用 $X_{ik},Y_{kj}$ 简化表示矩阵乘法 $X \cdot Y$ ：
$X_{ik}Y_{kj}:=sum_kX_{ik}Y_{kj}=O_{ij}$
其中 $O_{ij}$ 为输出矩阵的第 $i j$ 个元素。

求和简记法能以一种统一的方式表示各种张量运算，包括内积、外积、转置、点乘、矩阵的迹及其他自定义运算，为不同运算实现提供了统一模型.numpy 和 pytorch 中对其进行了实现，两者语义一致。

Torch 矩阵乘法：

print(a_tensor)
 
tensor([[11, 12, 13, 14],
        [21, 22, 23, 24],
        [31, 32, 33, 34],
        [41, 42, 43, 44]])
 
print(b_tensor)
 
tensor([[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3],
        [4, 4, 4, 4]])
 
# 'ik, kj -> ij'语义解释如下：
# 输入a_tensor: 2维数组，下标为ik,
# 输入b_tensor: 2维数组，下标为kj,
# 输出output：2维数组，下标为ij。
# 隐含语义：输入a,b下标中相同的k，是求和的下标，对应上面的例子2的公式
output = torch.einsum('ik, kj -> ij', a_tensor, b_tensor)
 
print(output)
 
tensor([[130, 130, 130, 130],
        [230, 230, 230, 230],
        [330, 330, 330, 330],
        [430, 430, 430, 430]])

高阶张量：
```
a = np.arange(60.).reshape(3,4,5)
b = np.arange(24.).reshape(4,3,2)
 
# 语义解析：
# 输入a：3阶张量，下标为ijk
# 输入b: 3阶张量，下标为jil
# 输出o: 2阶张量，下标为k和l
# 隐含语义：对i,j进行求和，公式在后面
o = np.einsum('ijk,jil->kl', a, b)
print(o)
 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])
 
# 验证：
print(np.sum(a[:,:,0]*b[:,:,0].T))
 
4400.0
 
print(np.sum(a[:,:,1]*b[:,:,0].T))
 
4532.0
```
公式为：
$A_{ijk}B_{jil}:=sum_isum_jA_{ijk}B_{jil}=O_{kl}$
$O$ 第 $k, l$ 个元素为矩阵 A[:,:,k] 和矩阵 B[:,:,l] 转置，对应元素相乘再求和。