我上一篇博文的一位讀者向我指出，對於切片matmul 之類的操作，np.einsum 比np.matmul 慢得多，除非您在參數(shù)列表中打開優(yōu)化標(biāo)誌： np.einsum(. . ., 最佳化= True).

帶著一些懷疑，我啟動了 Jupyter 筆記本並做了一些初步測試。 我天哪，這是完全正確的 - 即使對於兩個操作數(shù)的情況，優(yōu)化根本不應(yīng)該產(chǎn)生任何區(qū)別！

檢定 1 非常簡單 - 兩個不同維度的 C 階（又稱行主階）矩陣的矩陣乘法。 np.matmul 總是快二十倍左右。

對於檢定2，當(dāng)optimize=True時，結(jié)果截然不同。 np.einsum 仍然較慢，但最壞情況下僅慢 1.5 倍左右！

為什麼？

我對最佳化標(biāo)誌的理解是，當(dāng)存在三個或更多操作數(shù)時，它確定最佳收縮順序。 這裡，我們只有兩個操作數(shù)。 所以優(yōu)化應(yīng)該不會有什麼不同，對吧？

但也許優(yōu)化不只是選擇收縮順序？ 也許優(yōu)化器知道記憶體佈局，這與行優(yōu)先與列優(yōu)先佈局有關(guān)？

在矩陣乘法的小學(xué)方法中，要計算單個條目，您將迭代op1 中的行，同時迭代op2 中的列，因此將第二個參數(shù)按列優(yōu)先順序放置可能會導(dǎo)致加速對於np.einsum （假設(shè)np.einsum 有點像是底層矩陣乘法的小學(xué)方法的通用版本，我懷疑這是真的）。

因此，對於 測試 3，我為第二個運算元傳遞了一個列主矩陣，以查看當(dāng) optimize=False 時這是否會加快 np.einsum 的速度。

這是結(jié)果。 令人驚訝的是，np.einsum 還是相當(dāng)更糟。 顯然，發(fā)生了一些我不明白的事情 - 當(dāng) optimize 為 True 時，也許 np.einsum 使用完全不同的程式碼路徑？ 是時候開始挖掘了。

更深入

Numpy 1.12.0 的發(fā)行說明提到了最佳化標(biāo)誌的引入。 然而，優(yōu)化的目的似乎是確定操作數(shù)鏈中參數(shù)的組合順序（即關(guān)聯(lián)性） - 因此優(yōu)化不應(yīng)僅對兩個操作數(shù)產(chǎn)生影響，對吧？ 以下是發(fā)行說明：

np.einsum 現(xiàn)在支援最佳化參數(shù)，它將最佳化收縮順序。例如，np.einsum 將在一次傳遞中完成鏈點範(fàn)例np.einsum('ij,jk,kl->il', a, b, c)，其縮放比例類似於N^4；然而，當(dāng)optimize= True時，np.einsum將建立一個中間數(shù)組，以將此縮放減少到N^3或有效地np.dot(a, b).dot(c)。使用中間張量來減少縮放已應(yīng)用於通用 einsum 求和符號。有關(guān)更多詳細(xì)信息，請參閱 np.einsum_path。

為了讓這個謎團更加複雜，一些後來的發(fā)行說明表明 np.einsum 已升級為使用tensordot（它本身在適當(dāng)?shù)那闆r下使用BLAS）。 現(xiàn)在，這似乎很有希望。

但是，為什麼我們只在最佳化為True時看到加速？ 發(fā)生什麼事了？

如果我們在numpy/numpy/_core/einsumfunc.py 中閱讀def einsum(*operands, out=None, optimization=False, **kwargs) ，我們幾乎會立即看到這個提前退出的邏輯：

c_einsum 是否使用tensordot？ 我對此表示懷疑。 稍後在程式碼中，我們看到 1.14 註解似乎引用了 tensordot 呼叫：

所以，這就是發(fā)生的事情：

1. 如果 optimization 為 True，則收縮清單循環(huán)將被執(zhí)行 - 即使在簡單的兩個操作數(shù)情況下也是如此。
2. tensordot僅在收縮_列表循環(huán)中呼叫。
3. 因此，當(dāng)optimize為True時，我們呼叫tensordot（因此也呼叫BLAS）。

對我來說，這似乎是個錯誤。 恕我直言， np.einsum 開頭的「提前退出」仍應(yīng)偵測運算元是否與tensordot相容，並在可能的情況下呼叫tensordot。 然後，即使優(yōu)化為 False，我們也會得到明顯的 BLAS 加速。 畢竟，最佳化的語義與收縮順序有關(guān)，而不是與 BLAS 的使用有關(guān)，我認(rèn)為這應(yīng)該是給定的。

這裡的好處是，調(diào)用 np.einsum 進行相當(dāng)於張量調(diào)用的操作的人將獲得適當(dāng)?shù)募铀伲瑥男阅芙嵌葋砜?，這使得 np.einsum 的危險性降低了一些。

c_einsum 實際上是如何運作的？

我深入研究了一些 C 程式碼來檢查它。 實現(xiàn)的核心就在這裡。

經(jīng)過大量的參數(shù)解析和參數(shù)準(zhǔn)備，確定了軸迭代順序並準(zhǔn)備了專用迭代器。 迭代器的每個收益都代表了同時跨越所有操作數(shù)的不同方式。

假設(shè)某些特殊情況最佳化不適用，則根據(jù)涉及的資料類型確定適當(dāng)?shù)某朔e和 (sop) 函數(shù)：

然後，在從迭代器返回的每個多操作數(shù)步幅上呼叫此乘積和 (sop) 運算，如下所示：

這就是我對 einsum 工作原理的理解，誠然，它仍然有點薄弱——它確實值得我花更多的時間來了解它。

但它確實證實了我的懷疑，它的作用就像小學(xué)矩陣乘法方法的廣義、千兆腦版本。 最終，它委託給一系列「乘積之和」運算，這些運算依賴於在操作數(shù)中移動的「跨步器」——與學(xué)習(xí)矩陣乘法時用手指所做的沒有太大不同。

概括

那為什麼當(dāng)你使用optimize=True呼叫np.einsum時通常會比較快呢？ 原因有兩個。

第一個（也是最初的）原因是它試圖找到最佳的收縮路徑。 然而，正如我所指出的，當(dāng)我們只有兩個操作數(shù)時，這應(yīng)該不重要，就像我們在效能測試中所做的那樣。

第二個（也是較新的）原因是，當(dāng)optimize=True時，即使在兩個運算元的情況下，它也會啟動一個程式碼路徑，在可能的情況下呼叫tensordot，而tensordot又會嘗試使用BLAS。 BLAS 與矩陣乘法一樣最佳化！

所以，雙重操作數(shù)加速之謎解決了！ 然而，我們還沒有真正涵蓋由於收縮順序而導(dǎo)致的加速特性。 這得等以後的貼文了！ 敬請期待！

以上是研究 np.einsum 的性能的詳細(xì)內(nèi)容。更多資訊請關(guān)注PHP中文網(wǎng)其他相關(guān)文章！