本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯(lián)系出處。

那個在GitHub標星4.3k的張量操作工具Einops，在開源三年后終于中了頂會！

這是一個統(tǒng)一的、通用的操作張量結構的方法，基于愛因斯坦求和約定（Einstein summation convention）的思路開發(fā)，能夠大幅提高代碼的可讀性和易修改性。

同時，Einops支持Pytorch、TensorFlow、Chainer、Jax、Gluon等多個深度學習框架，以及Numpy、Cupy等張量計算框架。

ICLR 2022將其接收為Oral論文的消息一出，無數(shù)白嫖多年的“精神股東”們紛紛奔走相告，認為Einops確實“當之無愧”：

不過，這一早就聲名在外，還有無數(shù)大牛站臺的工具，在投往頂會后卻并非一帆風順。

比愛因斯坦求和約定更好用的標記法

我們先來了解一下Einops的基本原理。

它的設計思路來自于愛因斯坦在1916年提出的愛因斯坦求和約定，也叫愛因斯坦標記法（Einstein notation）。

這一方法的規(guī)定是：當一組乘積中，有兩個變量的腳標一樣，就要對相同的兩個腳標求和。例如下圖中的aibi：

這樣書寫的好處是，避免公式里出現(xiàn)大量的求和符號，看起來更簡潔。

Numpy里的Einsum就是一種模仿愛因斯坦求和約定的方法，可以說，這種思路已經被廣泛使用。

而Einops正是基于Einsum進行了諸多改進，針對張量操作過程中一些以前難以解決的問題，提供了更加便利的方案。

比如，當僅通過Pytorch，以及結合Einops兩種方法來實現(xiàn)超分辨率（Super-resolution）時，后者無疑極大降低了代碼冗余：

△上：原版 下：結合Einops

可以看到，結合了Einops之后的代碼不需要特殊指令PixelShuffle，并且，還使用了模塊調用nn.ReLU以及nn.Sequential。

最后得到的輸出結果既不包含虛假坐標，還能在框架之間進行轉移。

為什么Einops能做到這些？

這得益于Einops的本質：這是一種針對變換模式的新的標記法，能夠確保元素在張量中的位置與坐標變量的值一對一映射。

△Numpy和Einops操作之間的對應關系

對比愛因斯坦求和約定（Einsum），Einops有幾個額外的特征：

• 減少僅存在于輸入中的坐標（例如可以使用max-reduction法）
• 重復僅存在于輸出中的坐標（張量值對于新坐標的索引都是一樣的）
• 使表達式兩邊的所有坐標的標記唯一（Einsum允許重復）

在這些特征中，輸入和輸出被描述為張量的維度和坐標的預期順序，這使得基于Einops的代碼的可讀性和易修改性非常高，用戶也不需要在每次操作后記住或推斷出張量的形狀。

同時，Einops將輸入坐標（或其組成）與輸出坐標連接起來，這也使得張量結構在設計上就無法被破壞。

用括號表示的坐標的組成和分解也是Einops的一個主要創(chuàng)新之處：

通過上述的特點，Einops模式有利于更加靈活地處理高維度數(shù)據(jù)。

比如，一個注意力機制函數(shù)接受了一個形狀為[batch，seq，channel]的張量k q v，Einops可以將其高度、寬度和深度合成為一個維度。

同時，張量的頭部和batch維度也能夠被分組，這就保證了注意力頭部的獨立處理，從而將其變?yōu)槿S數(shù)據(jù)的多頭注意力：

類似的，通過修改輸入和輸出的結構，其他神經塊也能夠在Einops模式下完成“升級”。

論文接收并非一帆風順

Einops的作者是來自俄羅斯的Alex Rogozhnikov，他擁有莫斯科大學的數(shù)學和物理學博士學位，目前的主要研究領域是機器學習。

而Einops正是他最火熱的一項開發(fā)，不僅在GitHub上收獲4.3k星，還有特斯拉AI高管、FAIR實驗室工程師站臺：

但是，當作者將其整理成論文并投稿頂會時，有不少審稿人認為這篇論文“讀起來像是一篇技術博客”、“創(chuàng)新性不夠”、“缺乏嚴謹性”，并給出了3分（拒絕）的意見：

不過，會議主席最終給出了一錘定音的正面評價，并同意將其接收為Oral論文。

比起將幾個模塊共同訓練、壓縮、結合，最后在某個任意基準上的SOTA增加0.31 +/-1.04，這篇論文所描述的技術對于ICLR讀者來說更重要。

有Einops的使用者專門將會議主席最后的評價貼了出來，并表示：

除了新穎的技術和SOTA之外，其他的很多論文也同樣具有價值。

而不管這篇論文有何更廣泛的鼓勵意義，至少對于開發(fā)者Alex Rogozhnikov本人來說，也算是修成正果了。

?論文：

??https://openreview.net/forum?id=oapKSVM2bcj??

GitHub鏈接：

??https://github.com/arogozhnikov/Einops??

作者主頁：

??http://arogozhnikov.github.io/about/??