Numpy庫、pytorch庫和tensorflow庫中，存在一個非常通用但鮮為人知的函數(shù)，稱為einsum()，它根據(jù)愛因斯坦的求和約定執(zhí)行求和。PyTorch和TensorFlow像numpy支持einsum的好處之一是einsum可以用于神經(jīng)網(wǎng)絡(luò)架構(gòu)的任意計算圖，并且可以反向傳播。這是一個高效的符號計算，用于矩陣的各種求和操作， 在本教程文章中，我們揭開了einsum()的神秘面紗。
學(xué)習(xí)它的目的
愛因斯坦求和提供了一種緊湊而優(yōu)雅的方式來指定幾乎任何標(biāo)量/矢量/矩陣/張量的求和操作。 非常普遍，又減少計算機科學(xué)家所犯錯誤的數(shù)量，并減少他們花在推理線性代數(shù)上的時間。 通過同時更清晰，更明確，更自我表示，更具說明性和更少認(rèn)知負(fù)擔(dān)來實現(xiàn)。 它比矩陣乘法這樣的優(yōu)點在于它讓用戶不必考慮：
提供參數(shù)張量的正確順序
適用于參數(shù)張量的正確換位
確保正確的張量尺寸相互排列
-
正確的換位應(yīng)用于結(jié)果張量
愛因斯坦求和確實以著名的物理學(xué)家和理論家阿爾伯特愛因斯坦的名字命名。 但是，愛因斯坦沒有參與其發(fā)展。 他只是通過表達(dá)式的來推廣它。 在給Tullio Levi-Civita的一封信中，共同開發(fā)人員與Ricci演算的Gregorio Ricci-Curbastro一起（其求和符號只是其中的一部分），愛因斯坦寫道：I admire the elegance of your method of computation; it must be nice to ride through these fields upon the horse of true mathematics while the like of us have to make our way laboriously on foot.
愛因斯坦本人也高度贊揚這個符號求值
NB: As a further aside, the most general formulation of Einstein summation involves topics such as covariance and contravariance, indicated by subscript and superscript indices respectively. For our purposes, we will ignore co-/contravariance, since we can and will choose the “basis” we operate in to make the complexities that they introduce disappear.
einsum工作原理
一旦掌握愛因斯坦求和 公式你會非常方便使用。
導(dǎo)入庫
import tensorflow as tf
import numpy as np
import torch
它使用格式字符串和任意數(shù)量的參數(shù)張量，并返回結(jié)果張量。
格式化字符串語法：
- 逗號表示分隔參數(shù)，參數(shù)規(guī)范的數(shù)量和參數(shù)需要匹配
- 結(jié)果和參數(shù)的分割使用箭頭，箭頭是必須有的
- 參數(shù)和結(jié)果張量的規(guī)范是一系列（字母，ASCII）字符
- 張量規(guī)格中的字符數(shù)正好等于此張量的維數(shù)。
示例如下：
v = np.arange(100)
M = np.arange(16).reshape(4,4)
A = np.arange(25).reshape(5,5)
B = np.arange(20).reshape(5,4)
s = np.einsum('a->', v)
T = np.einsum('ij->ji', M)
C = np.einsum('mn,np->mp', A,B)
assert v.ndim == len('a')
assert s.ndim == len('')
assert M.ndim == len('ij')
assert T.ndim == len('ji')
assert A.ndim == len('mn')
assert B.ndim == len('np')
assert C.ndim == len('mp')
工作機理
內(nèi)部工作
在愛因斯坦求和Numpy.einsum()中，用一個字母標(biāo)記每個張量的每個軸，該字母表示在該軸上迭代時將使用的索引。 然后，einsum()很容易表達(dá)為一組深層嵌套的for循環(huán)。 這些for循環(huán)的核心是參數(shù)乘積的總和。例子如下
矩陣轉(zhuǎn)置
import torch
a = torch.arange(24).reshape(4, 6)
torch.einsum('ij->ji', [a])
tensor([[ 0, 6, 12, 18],
[ 1, 7, 13, 19],
[ 2, 8, 14, 20],
[ 3, 9, 15, 21],
[ 4, 10, 16, 22],
[ 5, 11, 17, 23]])
矩陣求和
a = torch.arange(6).reshape(2, 3)
torch.einsum('ij->', [a])
tensor(15)
矩陣列求和
a = torch.arange(6).reshape(2, 3)
torch.einsum('ij->j', [a])
tensor([3, 5, 7])
矩陣行求和
a = torch.arange(6).reshape(2, 3)
b = torch.arange(3)
torch.einsum('ik,k->i', [a, b])
tensor([ 5, 14])
點乘
a = torch.arange(6).reshape(2, 3)
b = torch.arange(15).reshape(3, 5)
torch.einsum('ik,kj->ij', [a, b])
tensor([[ 25, 28, 31, 34, 37],
[ 70, 82, 94, 106, 118]])
2D矩陣抽取
a = torch.arange(9).reshape(3, 3)
torch.einsum('ii->i', a)
tensor([0, 4, 8])
2D矩陣跡
a = torch.arange(9).reshape(3, 3)
torch.einsum('ii->', a)
tensor(12)
二次形式
batch矩陣相乘
a = torch.randn(3,2,5)
b = torch.randn(3,5,3)
torch.einsum('ijk,ikl->ijl', [a, b])
點積
a = torch.arange(3)
b = torch.arange(3,6) # [3, 4, 5]
torch.einsum('i,i->', [a, b])
一個MPL示例
# 15: MLP Backprop done easily (stochastic version).
# h = sigmoid(Wx + b)
# y = softmax(Vh + c)
Ni = 784
Nh = 500
No = 10
W = np.random.normal(size = (Nh,Ni)) # Nh x Ni
b = np.random.normal(size = (Nh,)) # Nh
V = np.random.normal(size = (No,Nh)) # No x Nh
c = np.random.normal(size = (No,)) # No
# Load x and t...
x, t = train_set[k]
# With a judicious, consistent choice of index labels, we can
# express fprop() and bprop() extremely tersely; No thought
# needs to be given about the details of shoehorning matrices
# into np.dot(), such as the exact argument order and the
# required transpositions.
#
# Let
#
# 'i' be the input dimension label.
# 'h' be the hidden dimension label.
# 'o' be the output dimension label.
#
# Then
# Fprop
ha = np.einsum("hi, i -> h", W, x) + b
h = sigmoid(ha)
ya = np.einsum("oh, h -> o", V, h) + c
y = softmax(ya)
# Bprop
dLdya = y - t
dLdV = np.einsum("h , o -> oh", h, dLdya)
dLdc = dLdya
dLdh = np.einsum("oh, o -> h ", V, dLdya)
dLdha = dLdh * sigmoidgrad(ha)
dLdW = np.einsum("i, h -> hi", x, dLdha)
dLdb = dLdha
TreeQN
我曾經(jīng)在實現(xiàn)TreeQN（ arXiv:1710.11417）的等式6時使用了einsum：給定網(wǎng)絡(luò)層l上的低維狀態(tài)表示zl，和激活a上的轉(zhuǎn)換函數(shù)Wa，我們想要計算殘差鏈接的下一層狀態(tài)表示。
在實踐中，我們想要高效地計算大小為B的batch中的K維狀態(tài)表示Z ∈ ?B × K，并同時計算所有轉(zhuǎn)換函數(shù)（即，所有激活A(yù)）。我們可以將這些轉(zhuǎn)換函數(shù)安排為一個張量W ∈ ?A × K × K，并使用einsum高效地計算下一層狀態(tài)表示。
import torch.nn.functional as F
def random_tensors(shape, num=1, requires_grad=False):
tensors = [torch.randn(shape, requires_grad=requires_grad) for i in range(0, num)]
return tensors[0] if num == 1 else tensors
# 參數(shù)
# -- [激活數(shù) x 隱藏層維度]
b = random_tensors([5, 3], requires_grad=True)
# -- [激活數(shù) x 隱藏層維度 x 隱藏層維度]
W = random_tensors([5, 3, 3], requires_grad=True)
def transition(zl):
# -- [batch大小 x 激活數(shù) x 隱藏層維度]
return zl.unsqueeze(1) + F.tanh(torch.einsum("bk,aki->bai", [zl, W]) + b)
# 隨機取樣仿造輸入
# -- [batch大小 x 隱藏層維度]
zl = random_tensors([2, 3])
transition(zl)
注意力
# 參數(shù)
# -- [隱藏層維度]
bM, br, w = random_tensors([7], num=3, requires_grad=True)
# -- [隱藏層維度 x 隱藏層維度]
WY, Wh, Wr, Wt = random_tensors([7, 7], num=4, requires_grad=True)
# 注意力機制的單次應(yīng)用
def attention(Y, ht, rt1):
# -- [batch大小 x 隱藏層維度]
tmp = torch.einsum("ik,kl->il", [ht, Wh]) + torch.einsum("ik,kl->il", [rt1, Wr])
Mt = F.tanh(torch.einsum("ijk,kl->ijl", [Y, WY]) + tmp.unsqueeze(1).expand_as(Y) + bM)
# -- [batch大小 x 序列長度]
at = F.softmax(torch.einsum("ijk,k->ij", [Mt, w]))
# -- [batch大小 x 隱藏層維度]
rt = torch.einsum("ijk,ij->ik", [Y, at]) + F.tanh(torch.einsum("ij,jk->ik", [rt1, Wt]) + br)
# -- [batch大小 x 隱藏層維度], [batch大小 x 序列維度]
return rt, at
# 取樣仿造輸入
# -- [batch大小 x 序列長度 x 隱藏層維度]
Y = random_tensors([3, 5, 7])
# -- [batch大小 x 隱藏層維度]
ht, rt1 = random_tensors([3, 7], num=2)
rt, at = attention(Y, ht, rt1)
