【深度学习】爱因斯坦求和约定(einsum)

import tensorflow as tf
print(tf.__version__)

2.0.0

一、爱因斯坦求和约定(einsum)的介绍

爱因斯坦求和约定是一种对复杂张量运算的优雅表达方式。在实现深度学习模型时，使用爱因斯坦求和约定可以编写更加紧凑和高效的代码。

einsum省略求和符号并隐式累加重复下标和输出未指明的下标。例如将两个矩阵$A\in {\mathbb{R}}^{I\times K}$和$B\in {\mathbb{R}}^{K\times J}$相乘，接着计算每列的和，最终得到向量$c\in {\mathbb{R}}^J$。使用einsum就可以表示为：

$c_j={\sum }_i{\sum }_k{A}_{ik}{B}_{kj}=A_{ik}{B}_{kj}$

省略掉中间带求和符号的表达式就是einsum表达式，其含义是计算行向量$A_{i,:}$按位乘以列向量$B_{:,j}$，然后求和(由于重复下标k所以求和，这里相当于点积)。这时得到的是两矩阵相乘后的矩阵$C\in {\mathbb{R}}^{I\times J}$，但是由于表达式指明为$c_j$，因此需要将计算得到的矩阵中的i维度进行求和(因为$c_j$中没有i)。这就实现了上面复杂的张量运算，下面是使用tensorflow计算上面的例子。

A = tf.ones((2,3))
B = tf.ones((3,4))*2
c = tf.einsum("ik,kj->j",A,B)
print(c)

tf.Tensor([12. 12. 12. 12.], shape=(4,), dtype=float32)

二、使用einsum实现深度学习中的常见操作

1.矩阵转置

A = tf.random.uniform((3,4))
B = tf.einsum("ij->ji",A)
print(A.shape)
print(B.shape)

(3, 4)
(4, 3)

2.求和

A = tf.random.uniform((3,4))
print(tf.einsum("ij->",A)) # 所有元素求和
print(tf.einsum("ij->j",A).shape) # 列求和
print(tf.einsum("ij->i",A).shape) # 行求和

tf.Tensor(6.1462874, shape=(), dtype=float32)
(4,)
(3,)

3.乘积

A = tf.random.uniform((2,3))
B = tf.random.uniform((3,4))
print(tf.einsum("ik,kj->ij",A,B).shape) # 矩阵乘法
a = tf.random.uniform((3,))
b = tf.random.uniform((3,))
print(tf.einsum("i,i->",a,b).shape) # 内积
print(tf.einsum("i,j->ij",a,b).shape) # 外积

(2, 4)
()
(3, 3)

4.batch乘法

在深度学习中样本通常是成批次的送到模型中，因此往往需要对每个样本对应的矩阵进行乘法，而由于整个tensor的第1个维度为batch size，因此要实现batch乘法会比较复杂，而使用einsum就非常的简洁。

A = tf.random.uniform((32,128,100))
B = tf.random.uniform((32,100,50))
print(tf.einsum("bsd,bdi->bsi",A,B).shape)

(32, 128, 50)

5.张量缩约

将两个高维张量在某些维度相乘，并在这些维度上求和

A = tf.random.uniform((3,5,7,2,8))
B = tf.random.uniform((4,7,2,3,5))
print(tf.einsum("pqrst,urspq->pqrstu",A,B).shape)

(3, 5, 7, 2, 8, 4)

参考