看看愛因斯坦對數(shù)據(jù)科學的小貢獻

怎么，愛因斯坦（Albert Einstein）那會兒就有數(shù)據(jù)科學了嗎？

倒不是這個意思，愛因斯坦也沒有提出什么數(shù)學理論，但他提出了一個針對數(shù)學公式的符號簡化辦法，即愛因斯坦求和約定（Einstein Summation Convention）或者叫愛因斯坦標記法（Einstein Notation）。

他這套方法不僅方便了相關理論的書寫，而且意外地給如今數(shù)據(jù)科學中編程實現(xiàn)相關計算帶來了方便，讓我們來看看到底怎么回事。

1引言

先瞄一眼愛因斯坦當年在論文廣義相對論的基礎里涉及的一些數(shù)學公式。

看著比較繁瑣吧，如果學過微分幾何的應該熟悉這套的。

意大利數(shù)學家雷戈里奧·里奇（Gregorio Ricci-Curbastro）在研究黎曼、克里斯托費爾等人的黎曼幾何微分不變量時提出了絕對微分學，然后與他學生勒維·季維塔（Levi-Givita）創(chuàng)立了張量分析。而黎曼幾何和張量分析正是愛因斯坦廣義相對論的數(shù)學基礎。

要注意的是，張量分析中的張量指一種特殊的多重線性映射，在相關基底給定的情況下可以用一個多維數(shù)組表示，而 NumPy 和 PyTorch 等軟件包中的張量特指多維數(shù)組。

在張量分析中，為了分別處理張量隨坐標變換的協(xié)變和逆變，引入了上下標，上標表示逆變張量的分量，下標表示協(xié)變張量的分量，它們根據(jù)基底的變化分別進行逆變或協(xié)變。

通常，當處理協(xié)變張量和逆變張量時，其中上下標的位置也指示了張量類型以及縮并方式。

但是，愛因斯坦求和約定也可以有不同應用方式。例如，更一般的情況是坐標基底固定，或者不考慮坐標的情況時，則可以選擇僅使用下標。在 NumPy 或 PyTorch 等程序中涉及的更多情況就是這個樣子。

像向量內(nèi)積，矩陣-向量乘法，以及矩陣乘法都可以用這套標記法來簡化書寫。

例如向量內(nèi)積，

矩陣 乘以向量 ，

可以記為，

而矩陣 和 相乘，

可以記為，

2NumPy 之 einsum

本文主要介紹 NumPy 中實現(xiàn)的愛因斯坦求和約定，而像 PyTorch、TensorFlow 等情況基本類似。

這里主要拿來處理高維數(shù)組的相關計算，并不需要上下標那一套，而是統(tǒng)一針對下標來簡化求和連加符。使用時需要明確的是下面這幾點，

• 哪幾個輸入數(shù)組參與運算
• 沿輸入的哪些軸（維度）計算乘積
• 輸出數(shù)組需要保留哪些軸

具體反映在函數(shù) np.einsum 的參數(shù)上，比如下圖所示的例子，三個顏色對應三個輸入張量，分別是 2 維數(shù)組，3 維數(shù)組和 2 維數(shù)組，而輸出張量是 2 維數(shù)組。

看個例子

為了快速了解 einsum，我們直接來看一個向量乘以矩陣的例子。

import?numpy?as?np

A?=?np.array([0,1,2])
A

array([0, 1, 2])

B?=?np.arange(12).reshape(3,?4)
B

array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

將 A 看成行向量，然后乘以（矩陣乘法）B，得到一個 1x4 的向量。

#?這樣嗎？
A*B

我們知道，NumPy 這里的運算是元素級別的運算，因此是通過廣播機制將 A 擴展為與 B 同樣大小的二維數(shù)組。但是，很可惜，此處 A 的 shape 為 (3)，而 B 的 shape 為 (3,4)，它們并不符合廣播的條件，因此出錯。

ValueError: operands could not be broadcast together with shapes (3,) (3,4)

增加一個新軸，讓 A 的 shape 變成 (3,1)，然后就可以廣播了。如果不清楚?NumPy 的廣播機制，可以出門左拐看另一篇哦。

C?=?A[:,?np.newaxis]*B?
C

array([[ 0,  0,  0,  0],
       [ 4,  5,  6,  7],
       [16, 18, 20, 22]])

然后將每一列的數(shù)字加起來，得到一個有 4 個元素的向量。

C.sum(axis=0)

array([20, 23, 26, 29])

#?組合在一起
(A[:,?np.newaxis]*?B).sum(axis=0)

array([20, 23, 26, 29])

使用愛因斯坦求和約定來實現(xiàn)，將變得更加簡潔高效。

D?=?np.einsum('i,ij->j',?A,?B)
D

array([20, 23, 26, 29])

字符串 'i,ij->j' 由 '->' 分成了兩部分，它左邊的 'i,ij' 對應兩個輸入，而右邊的 'j' 對應輸出。輸出中沒有下標 'i'，說明對兩個輸入沿著這個下標求和，而 j 所在的軸仍然保留。而 j 下標有 0 到 3 共 4 個值，因此最終得到一個有 4 個元素的向量。對應如下公式，

對比一下性能，

%timeit?(A[:,?np.newaxis]*?B).sum(axis=0)

3.46 μs ± 429 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit?np.einsum('i,ij->j',?A,?B)

1.18 μs ± 34.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

當然，這個簡單例子我們自然可以用內(nèi)積 dot 來計算。這也是為數(shù)不多的性能比愛因斯坦求和約定更好的情況，而大多數(shù)情況下，愛因斯坦求和約定的性能更優(yōu)。

%timeit?A.dot(B)

655 ns ± 12.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

#?或者
np.dot(B.T,A)

array([20, 23, 26, 29])

但是，einsum 可以很方便地實現(xiàn)如下計算，

E?=?np.einsum('i,ij->i',?A,?B)
E

array([ 0, 22, 76])

此時，輸出中沒有下標 'j'，說明對兩個輸入沿著這個下標求和，而 i 所在的軸仍然保留。而 i 下標有 0 到 2 共 3 個值，因此最終得到一個有 3 個元素的向量。對應如下公式，

注意，連加符里的下標會消失，但沒有出現(xiàn)在連加符里的下標以及相應的軸會保留。

使用說明

一般來說，使用 einsum 時是為了對輸入的一些數(shù)組沿著某些軸作乘積運算，那對這些軸當然有一定要求，比如沿著相同標號的軸的元素個數(shù)一樣多。

具體操作時，不妨先寫出你要計算的數(shù)學公式，然后把連加符去掉，再根據(jù)輸入輸出的下標確定 einsum 中的參數(shù)。

下圖給出一個例子，先寫出計算公式，再轉(zhuǎn)化為 np.einsum 里的字符串。

關鍵是為輸入數(shù)組的軸和我們想要輸出的數(shù)組選擇正確的標簽。可以選擇兩種方式之一執(zhí)行此操作，使用字符串或者使用整數(shù)列表。這里使用前者，即使用字符串來表達數(shù)學公式。

一個很好的例子是矩陣乘法，它將行與列相乘，然后對乘積結(jié)果求和。

• 對于兩個二維數(shù)組 A 和 B，矩陣乘法操作可以用 np.einsum('ij,jk->ik', A, B) 完成。
• 字符串 'ij,jk->ik' 是什么意思？
• 它被箭頭 -> 分成兩部分。左側(cè)部分標記輸入數(shù)組的軸：'ij' 標記 A 以及 'jk' 標記 B。字符串的右側(cè)部分用字母 'ik' 標記單個輸出數(shù)組的軸。
• 即輸入兩個二維數(shù)組，輸出一個新的二維數(shù)組。

A?=?np.array([[1,1,1],
??????????????[2,2,2],
??????????????[5,5,5]])

B?=?np.array([[0,1,0],
??????????????[1,1,0],
??????????????[1,1,1]])

np.einsum('ij,jk->ik',?A,?B)

array([[ 2,  3,  1],
       [ 4,  6,  2],
       [10, 15,  5]])

看下圖，注意軸的顏色，j 所在的軸被縮并掉了。

真相，'ij,jk->ik' 背后對應的數(shù)學公式

np.einsum 可以替代如下常用的運算，

• 矩陣求跡: trace
• 求矩陣對角線: diag
• 張量（沿軸）求和: sum
• 張量轉(zhuǎn)置: transopose
• 矩陣乘法: dot
• 張量乘法: tensordot
• 向量內(nèi)積: inner
• 外積: outer

a?=?np.array([1,1,1])
b?=?np.array([2,3,4])

%%timeit?
np.inner(a,?b)

765 ns ± 13.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%%timeit?
np.einsum('i,i->',?a,?b)

1.06 μs ± 2.13 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%%time
np.outer(a,b)

CPU times: user 44 μs, sys: 0 ns, total: 44 μs
Wall time: 46.7 μs

array([[2, 3, 4],
       [2, 3, 4],
       [2, 3, 4]])

#?用?einsum?計算外積
np.einsum('i,j->ij',a,b)

array([[2, 3, 4],
       [2, 3, 4],
       [2, 3, 4]])

%%timeit?
np.einsum('i,j->ij',a,b)

1.18 μs ± 2.12 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

求跡

a?=?np.arange(9).reshape((3,3))
a

array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

np.trace(a)

12

np.einsum('ii->',a)

12

#?或者
np.einsum('ii',a)

12

矩陣乘法 dot

#?1d
a?=?np.array([1,1,1])
b?=?np.array([2,3,4])

%%time
#?dot
a.dot(b)

CPU times: user 20 μs, sys: 0 ns, total: 20 μs
Wall time: 23.4 μs

9

%%time
np.einsum('i,i->',?a,b)

CPU times: user 47 μs, sys: 1 μs, total: 48 μs
Wall time: 53.4 μs

9

#?2d
a?=?np.arange(20).reshape(4,5)
b?=?np.arange(15).reshape(5,3)

%%time
#?dot
a.dot(b)

CPU times: user 33 μs, sys: 0 ns, total: 33 μs
Wall time: 38.6 μs

array([[ 90, 100, 110],
       [240, 275, 310],
       [390, 450, 510],
       [540, 625, 710]])

%%time
a@b

CPU times: user 48 μs, sys: 1e+03 ns, total: 49 μs
Wall time: 53.9 μs

array([[ 90, 100, 110],
       [240, 275, 310],
       [390, 450, 510],
       [540, 625, 710]])

%%time
np.einsum('ij,jk->ik',?a,b)

CPU times: user 30 μs, sys: 0 ns, total: 30 μs
Wall time: 32.4 μs

array([[ 90, 100, 110],
       [240, 275, 310],
       [390, 450, 510],
       [540, 625, 710]])

張量 dot

• 方法 np.tensordot 是沿著軸指定的子數(shù)組做點乘操作，即沿著 axes 指出的軸求和，最終這些軸就消失了。從這個例子可以看出，這個方法雖然使用方便，但性能遠不如上面幾個。

%%time
np.tensordot(a,?b,?axes=([1,],[0,]))

CPU times: user 209 μs, sys: 4 μs, total: 213 μs
Wall time: 211 μs

array([[ 90, 100, 110],
       [240, 275, 310],
       [390, 450, 510],
       [540, 625, 710]])

思考，'ij,jk->ijk' 是在干嘛？

r?=?np.einsum('ij,jk->ijk',?a,?b)
r

array([[[  0,   0,   0],
        [  3,   4,   5],
        [ 12,  14,  16],
        [ 27,  30,  33],
        [ 48,  52,  56]],

       [[  0,   5,  10],
        [ 18,  24,  30],
        [ 42,  49,  56],
        [ 72,  80,  88],
        [108, 117, 126]],

       [[  0,  10,  20],
        [ 33,  44,  55],
        [ 72,  84,  96],
        [117, 130, 143],
        [168, 182, 196]],

       [[  0,  15,  30],
        [ 48,  64,  80],
        [102, 119, 136],
        [162, 180, 198],
        [228, 247, 266]]])

常用 einsum 操作

1 維數(shù)組

舉例

a?=?np.arange(9)
np.einsum('i',?a)

array([0, 1, 2, 3, 4, 5, 6, 7, 8])

np.einsum('i->',?a)

36

np.einsum('i,i',?a,?a)

204

np.einsum('i,i->',?a,?a)

204

r?=?np.einsum('i,j->ij',?a,?a)
r

array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4,  5,  6,  7,  8],
       [ 0,  2,  4,  6,  8, 10, 12, 14, 16],
       [ 0,  3,  6,  9, 12, 15, 18, 21, 24],
       [ 0,  4,  8, 12, 16, 20, 24, 28, 32],
       [ 0,  5, 10, 15, 20, 25, 30, 35, 40],
       [ 0,  6, 12, 18, 24, 30, 36, 42, 48],
       [ 0,  7, 14, 21, 28, 35, 42, 49, 56],
       [ 0,  8, 16, 24, 32, 40, 48, 56, 64]])

np.sum(r)

1296

np.einsum('i,j->',?a,?a)

1296

2 維數(shù)組

舉例

#?2d
A?=?np.arange(12).reshape(4,3)
B?=?np.arange(12).reshape(4,3)

np.einsum('ij,kj->ik',?A,?B)

array([[  5,  14,  23,  32],
       [ 14,  50,  86, 122],
       [ 23,  86, 149, 212],
       [ 32, 122, 212, 302]])

np.einsum('ij,kj->ikj',?A,?B).shape

(4, 4, 3)

np.einsum('ij,kl->ijkl',?A,?B).shape

(4, 3, 4, 3)

3不同方法?PK

張量運算看著挺復雜的，但也有跡可循，而且有很多方法來實現(xiàn)同一個運算。最后再來一個例子，比較一下張量運算的不同實現(xiàn)及其性能。

• 創(chuàng)建一維數(shù)組 a，[0,1,...,59]，將其轉(zhuǎn)化為 shape 為 (3,4,5)`` 的三維數(shù)組A`
• 創(chuàng)建一維數(shù)組 b，[0,1,...,23]，將其轉(zhuǎn)化為 shape 為 (4,3,2) 的三維數(shù)組 B
• 將 A 轉(zhuǎn)置成 shape 為 (4,3,5) 的數(shù)組 D
• 即用兩到三種方法計算如下公式:

A?=?np.arange(60.).reshape(3,4,5)
B?=?np.arange(24.).reshape(4,3,2)

方法一

直接使用 np.tensordot。

%%timeit
C?=?np.tensordot(A,?B,?axes=([1,0],[0,1]))

12.6 μs ± 409 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

方法二

利用 NumPy 的廣播機制。

D?=?np.transpose(A,[1,0,2])
D.shape

(4, 3, 5)

%%timeit
np.sum(D[...,None]*B[:,:,None,:],(0,1))

5.4 μs ± 43.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

方法三

使用 np.einsum。

%%timeit?
np.einsum('jik,ijl->kl',?A,?B)

2.66 μs ± 10.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

方法四

實際上，這個例子比較特殊，我們可以想一個辦法來進一步提高計算性能，那就是通過變形數(shù)組來使得可以用 dot 來代替該運算。而我們知道 dot 得到專門優(yōu)化，所以性能杠杠的。但并不是所有運算都這么幸運，不一定能這么轉(zhuǎn)化哦。

%%timeit?
A.T.reshape(A.shape[-1],?-1).dot(B.reshape(-1,?B.shape[-1]))

1.94 μs ± 13.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

注意，計算結(jié)果都是一樣的，

array([[4400.,?4730.],
???????[4532.,?4874.],
???????[4664.,?5018.],
???????[4796.,?5162.],
???????[4928.,?5306.]])

最后看下本文使用的 numpy 版本，不同版本會有不同表現(xiàn)哦。

np.version.version

'1.15.2'

本文通過一些例子展示了愛因斯坦求和約定的強大功能，那或許你會問: 在實際寫算法時有人用嗎？自然是有大量使用的，畢竟這個工具可以完成各種復雜計算，可謂張量（高維數(shù)組）計算大殺器，完全值得花點時間學習一番。限于篇幅，至于實際使用的具體例子，我們就留給下一篇吧。