前言

说明：讲解时会对相关文章资料进行思想、结构、优缺点，内容进行提炼和记录，相关引用会标明出处，引用之处如有侵权，烦请告知删除。
转载请注明：DengBoCong

这里我们来好好探讨一下深度学习中，矩阵乘法的使用，其实主要是围绕einsum来进行探讨，即通过爱因斯坦求和约定来指导矩阵乘法，同时附带陈列其他矩阵乘法的API，方便进行直观感受。本文中的计算框架及版本如下：

  • TensorFlow2.3
  • PyTorch1.7.0
  • Numpy1.19

爱因斯坦求和约定

我们讨论einsum绕不开爱因斯坦求和约定的，爱因斯坦求和约定（Einstein summation convention）是一种标记的约定，又称为爱因斯坦标记法（Einstein notation），在处理关于坐标的方程式时非常有用，用一句话来总结爱因斯坦求和约定，就是：

当式子中任何一个角标出现了两次，并且一次是上标、一次是下标时，那么该式表示的实际上是对这个角标一切可能值的求和。换言之，如果角标 $i$ 作为上标和下标各出现了一次，那么式子相当于添加了一个关于 $i$ 的求和符号

我们下面使用线性函数和矩阵运算来对爱因斯坦求和约定进行举例说明：

  • 线性函数：从张量中我们知道，一个1-线性函数可以表示为一个向量，这样的向量常被称为余向量、补向量或者1-形式。通常，我们用下标来表示一个余向量的各分量：$a=(a_1,a_2,a_3)$ ，而用上标来表示一个通常的几何向量：$v=(v^1,v^2,v^3)$。注意，上标不是乘方，则 $a$ 和 $v$ 的内积是：
    $$\sum_{i=1,2,3}^{}a_iv^i$$
    用爱因斯坦求和约定， $a$ 和 $v$ 的内积就可以写为 $a_iv^i$
  • 矩阵运算：对于矩阵 $A$，我们把其第 $i$ 行第 $j$ 列的元素表示为 $A_j^i$。则矩阵乘法表示为：如果 $A=BC$，那么 $A_j^i=B_k^iC_j^k$。矩阵 $A$ 的迹为 $A_i^i$

由于重复出现而实际上应该是求和的指标，被称为赝指标或者哑指标（dummy index），因为它们不是真正的指标，而是可以用任意字母代替的。没有求和的指标是固定的，是真正的指标．比如说 $B_k^iC_j^k$ 中 $k$ 可以是任何字母，但是 $i$ 和 $j$ 是不可以替换成别的字母的，因为它们由 $A_j^i$ 决定了。在这里，哑指标实际上是表示遍历全部可能的真指标。

爱因斯坦求和约定的表示方法脱胎于矩阵乘法的要求，但是却不依赖于矩阵行和列的形式，转而关注指标间的配合，相比传统的矩阵表达，能更方便地推广到高阶张量的情形中。本文关于爱因斯坦求和约定的相关点到为止，如果感兴趣其公式，可以参考这一篇文章：爱因斯坦求和约定。

einsum介绍

通过使用einsum函数，我们可以使用爱因斯坦求和约定（Einstein summation convention）在NumPy数组上指定操作。einsum函数由于其强大的表现力和智能循环，它在速度和内存效率方面通常可以超越我们常见的array函数。但缺点是，可能需要一段时间才能理解符号，有时需要尝试才能将其正确的应用于棘手的问题，当然熟悉之后得心应手才是使用关键。

einsum以一种优雅的方式，表示各种矩阵运算，好处在于你不需要去记和使用计算框架中（TensorFlow|PyTorch|Numpy）点积、外积、转置、矩阵-向量乘法、矩阵-矩阵乘法的函数名字和签名。从某种程度上解决引入不必要的张量变形或转置运算，以及可以省略的中间张量的现象。不仅如此，einsum有时可以编译到高性能代码，事实上，PyTorch最近引入的能够自动生成GPU代码并为特定输入尺寸自动调整代码的张量理解（Tensor Comprehensions）就基于类似einsum的领域特定语言。此外，可以使用opt einsum和tf einsum opt这样的项目优化einsum表达式的构造顺序。

假设我们有两个多维数组A和B，现在让我们要进行如下操作：

  • 以特定方式将A与B相乘以创建新的乘积数组
  • 沿特定轴求和该新数组
  • 以特定顺序转置新数组的轴

einsum帮助我们更快，更高效地执行此操作，当然，NumPy函数的组合（例如multiply，sum和transpose）也可以实现。

einsum应用

我们在这一节用Numpy的einsum进行讲解说明（Numpy中einsum先被开发出来，TensorFlow和PyTorch都在一定程度上参考了它），在下一节，我会将TensorFlow、PyTorch和Numpy的API都贴出来。

在Numpy中，einsum使用格式字符串和任意数量的Numpy张量作为参数调用，并返回结果张量。

在这里插入图片描述
格式字符串包含分隔参数说明的逗号（，）和将参数说明与张量的参数分开的箭头（->）。参数说明中的数量和参数的数量必须匹配，并且必须精确地出现一个箭头，后跟一个结果说明

在这里插入图片描述
格式字符串中的字符数完全等于该张量的维数。

在这里插入图片描述
下面展示einsum()格式字符串的完整示例。 尝试猜测结果将是什么？注意了，0维张量（标量）对于参数和结果均有效，由空字符串“”表示。 再次提醒您，格式字符串必须仅由ASCII小写或大写字母组成。

在这里插入图片描述
我们这里使用一个实际的例子来进行说明，首先我们要相乘的两个​​数组是：

1
2
3
4
5
6
A = array([[1, 1, 1],
           [2, 2, 2],
           [5, 5, 5]])
B = array([[0, 1, 0],
           [1, 1, 0],
           [1, 1, 1]])

我们的矩阵乘法np.einsum(‘ij,jk->ik’, A, B)大致如下：
在这里插入图片描述
要了解输出数组的计算方法，请记住以下三个规则：

  • 在输入数组中重复的字母意味着值沿这些轴相乘。乘积结果为输出数组的值。在本例中，我们使用字母 j 两次：A和B各一次。这意味着我们将A每一行与B每列相乘。这只在标记为 j 的轴在两个数组中的长度相同（或者任一数组长度为1）时才有效。

  • 输出中省略的字母意味着沿该轴的值将相加。在这里，j 不包含在输出数组的标签中。通过累加的方式将它从轴上除去，最终数组中的维数减少1。如果输出是’ijk’，我们得到的结果是3x3x3数组（如果我们不提供输出标签，只写箭头，则对整个数组求和）。

  • 我们可以按照我们喜欢的任何顺序返回未没进行累加的轴。如果我们省略箭头’->’，NumPy会将只出现一次的标签按照字母顺序排列（因此实际上’ij,jk->ik’相当于’ij,jk’）。如果我们想控制输出的样子，我们可以自己选择输出标签的顺序。例如，’ij,jk->ki’为矩阵乘法的转置。

现在，我们已经知道矩阵乘法是如何工作的。下图显示了如果我们不对 j 轴进行求和，而是通过写np.einsum(‘ij,jk->ijk’, A, B)将其包含在输出中，我们会得到什么。右边代表 j 轴已经求和：
在这里插入图片描述
注意，由于np.einsum(‘ij,jk->ik’, A, B)函数不构造3维数组然后求和，它只是将总和累加到2维数组中。下面是两个表格展示了einsum如何进行各种NumPy操作。我们可以用它来熟悉符号。

  • 让A和B是两个形状兼容的一维数组（也就是说，我们相应的轴的长度要么相等，要么其中一个长度为1）：
调用 Numpy等式 描述
(‘i’, A) A 返回数组A的视图
(‘i->’, A) sum(A) 数组A值的总和
(‘i,i->i’, A, B) A * B A和B的数组元素依次相乘
(‘i,i’, A, B) inner(A, B) A和B的点积（内积）
(‘i,j->ij’, A, B) outer(A, B) A和B的外积（叉积）
  • 现在，我们A和B是与之兼容形状的两个二维数组：
调用 Numpy等式 描述
(‘ij’, A) A 返回A的视图
(‘ji’, A) A.T A的转置
(‘ii->i’, A) diag(A) A的主对角线
(‘ii’, A) trace(A) A的主对角线的和
(‘ij->’, A) sum(A) A的值相加
(‘ij->j’, A) sum(A, axis=0) 通过A的轴竖直（列）求和
(‘ij->i’, A) sum(A, axis=1) 通过A的轴水平（行）求和
(‘ij,ij->ij’, A, B) A * B A和B逐元素依次相乘
(‘ij,ji->ij’, A, B) A * B.T A和B的转置逐个元素依次相乘
(‘ij,jk’, A, B) dot(A, B) A和B的矩阵乘法
(‘ij,kj->ik’, A, B) inner(A, B) A和B点积（内积）
(‘ij,kj->ijk’, A, B) A[:, None] * B A的每一行乘以B
(‘ij,kl->ijkl’, A, B) A[:, :, None, None] * B A的每个值乘以B

当处理大量维度时，别忘了einsum允许使用省略号语法’…’。这提供了一种变量的方式标记我们不大感兴趣的轴，例如np.einsum(‘…ij,ji->…’, a, b)，仅将a的最后两个轴与2维数组b相乘。

TensorFlow、PyTorch和Numpy

通常，要将元素式方程式转换为方程式字符串，可以使用以下过程（右侧是矩阵乘法示例的中间字符串）

  1. 原始元素式方程式：C[i,k] = sum_j A[i,j] * B[j,k]
  2. 删除变量名称，方括号和逗号：ik = sum_j ij * jk
  3. 将“*”替换成“,”：ik = ij, jk
  4. 去掉求和符号：ik = ij, jk
  5. 输出移到右边，将将“=”替换成“->”：ij,jk->ik
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
tf.einsum(equation, *inputs, **kwargs)
参数：
	equation：描述计算的字符串，格式与numpy.einsum相同。
	*inputs：输入（张量），其形状应与方程保持一致。
	**kwargs：
		optimize：用于使用opt_einsum查找计算路径的优化策略，可选项包括 'greedy', 'optimal', 'branch-2', 'branch-all' ,'auto'，默认是'greedy'
		name：操作名称（可选）
返回值：张量，形状与方程中一致
示例：
	# 矩阵相乘
	einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]
	# 点积
	einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]
	# 外积
	einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]
	# 转置
	einsum('ij->ji', m)  # output[j,i] = m[i,j]
	# 主对角线的和
	einsum('ii', m)  # output[j,i] = trace(m) = sum_i m[i, i]
	# 批量矩阵相乘
	einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]
1
2
3
4
5
6
7
8
9
10
11
torch.einsum(equation, *operands)
参数：
	equation：描述计算的字符串，格式与numpy.einsum相同。
	operands：输入（张量），其形状应与方程保持一致。
示例：
torch.einsum('i,j->ij', x, y)  # 外积
torch.einsum('bn,anm,bm->ba', l, A, r) # 计算torch.nn.functional.bilinear
torch.einsum('bij,bjk->bik', As, Bs) # 批量矩阵相乘
torch.einsum('ii->i', A) # 对角线之和
torch.einsum('...ii->...i', A) # 批量对角线之和
torch.einsum('...ij->...ji', A).shape # 批量转置
1
2
3
4
5
6
7
8
9
numpy.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False)
参数：
	subscripts：描述计算的字符串，除非包含显式指示符“->”以及精确输出形式的下标标签，否则将执行隐式（经典的爱因斯坦求和）计算
	operands：输入（张量），其形状应与方程保持一致。
	out：ndarray类型（可选），如果提供，则将计算结果放入此数组中
	dtype：{data-type, None}（可选），如果提供，则强制使用指定的数据类型计算。 请注意，你可能还必须提供一个更宽松的转换参数以允许进行转换。 默认为None。
	order：{‘C’, ‘F’, ‘A’, ‘K’}（可选），控制输出的内存布局。其中，‘C’表示C连续的。‘F’表示它应该是Fortran连续的。如果输入全为‘F’，则‘A’表示应为‘F’，否则为‘C’。‘K’表示它应尽可能与输入尽可能靠近布局，包括任意排列的轴。
	casting：{‘no’, ‘equiv’, ‘safe’, ‘same_kind’, ‘unsafe’}（可选），控制可能发生的数据类型转换。 不建议将其设置为“unsafe”，因为这会积聚不利影响。其中，“no”表示完全不应该转换数据类型，“equiv”表示仅允许按照字节顺序转换，“safe”表示只允许保留值的强制类型转换。“same_kind”表示仅允许安全类型转换或同一类型（例如float64到float32）内的类型转换。“unsafe”表示可能会进行任何数据转换。
	optimize：{False, True, ‘greedy’, ‘optimal’}（可选），控制优化策略，如果True将默认设置为“贪心”算法，如果是False则不会进行优化。 还接受np.einsum_path函数的显式收缩列表。

其他乘法

TensorFlow2.3

1
2
3
4
5
6
7
8
9
10
11
tf.linalg.matmul(a, b, transpose_a=False, transpose_b=False, adjoint_a=False, adjoint_b=False,a_is_sparse=False, b_is_sparse=False, name=None)
参数：
	a：类型为float16, float32, float64, int32, complex64, complex128，并且秩大于1
	b：和a相同类型和秩
	transpose_a：如果为True，a在计算前会被转置
	transpose_b：如果为True，b在计算前会被转置
	adjoint_a：如果为True，a在计算前会被共轭和转置
	adjoint_b：如果为True，b在计算前会被共轭和转置
	a_is_sparse：如果为True，则将a视为稀疏矩阵。
	b_is_sparse：如果为True，则将b视为稀疏矩阵。
	name：操作名称（可选）
1
2
3
4
5
6
7
tf.sparse.sparse_dense_matmul(sp_a, b, adjoint_a=False, adjoint_b=False, name=None)
参数：
	a：SparseTensor或者是密度矩阵，并且秩为2
	b：和a相同类型，密度矩阵或者是a：SparseTensor
	adjoint_a：如果为True，a在计算前会被共轭和转置
	adjoint_b：如果为True，b在计算前会被共轭和转置
	name：操作名称（可选）
  • tf.math.multiply
    1
    2
    3
    4
    5
    tf.math.multiply(x, y, name=None)
    参数：
    	a：类型为bfloat16, half, float32, float64, uint8, int8, uint16, int16, int32, int64, complex64, complex128的张量
    	b：和a相同类型
    	name：操作名称（可选）

Pytorch1.7.0

1
2
3
4
5
torch.matmul(input, other, *, out=None)
参数：
	input：第一个用于乘法的张量
	other：第二个用于乘法的张量
	out：可选，输出张量

Numpy1.19

1
2
3
4
5
numpy.matmul(x1, x2, /, out=None, *, casting='same_kind', order='K', dtype=None, subok=True[, signature, extobj])
参数：
	x1, x2：输入数组
	out：输出储存的位置，如果提供，则其形状必须与签名(n,k),(k,m)->(n,m)相匹配。 如果未提供或没有，则返回一个新分配的数组。
	**kwargs：对于其他仅关键字参数，参考文档

参考文献