【深度學習】愛因斯坦求和約定(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)

參考