import tensorflow as tf
print(tf.__version__)
2.0.0
一、爱因斯坦求和约定(einsum)的介绍
爱因斯坦求和约定是一种对复杂张量运算的优雅表达方式。在实现深度学习模型时,使用爱因斯坦求和约定可以编写更加紧凑和高效的代码。
einsum省略求和符号并隐式累加重复下标和输出未指明的下标。例如将两个矩阵 A ∈ R I × K A\in \mathbb{R}^{I\times K} A∈RI×K和 B ∈ R K × J B\in \mathbb{R}^{K\times J} B∈RK×J相乘,接着计算每列的和,最终得到向量 c ∈ R J c\in \mathbb{R}^J c∈RJ。使用einsum就可以表示为:
c j = ∑ i ∑ k A i k B k j = A i k B k j c_j = \sum_{i}\sum_{k}A_{ik}B_{kj}=A_{ik}B_{kj} cj=∑i∑kAikBkj=AikBkj
省略掉中间带求和符号的表达式就是einsum表达式,其含义是计算行向量 A i , : A_{i,:} Ai,:按位乘以列向量 B : , j B_{:,j} B:,j,然后求和(由于重复下标k所以求和,这里相当于点积)。这时得到的是两矩阵相乘后的矩阵 C ∈ R I × J C\in \mathbb{R}^{I \times J} C∈RI×J,但是由于表达式指明为 c j c_j cj,因此需要将计算得到的矩阵中的i维度进行求和(因为 c j c_j cj中没有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)