乾明 發自 凹非寺
量子位 報道 | 公眾號 QbitAI

AI框架，又來重磅中國玩家。

剛剛，清華自研的深度學習框架，正式對外開源?！百F系”計算機系的圖形實驗室出品，取名Jittor，中文名計圖。

值得一提的是，這也是首個來自中國高?？蒲袡C構的開源深度學習框架，之前，業內來自“高校”的還有加拿大蒙特利爾大學的Theano，UC伯克利的Caffe。

與主流的深度學習框架TensorFlow、Pytorch不同，Jittor是一個完全基于動態編譯（Just-in-time）、使用元算子和統一計算圖的深度學習框架。

研發團隊介紹稱，開發Jittor是為了將新技術、硬件和模型的能力，更好地釋放出來。

“深度學習發展迅猛，TensorFlow、PyTorch這些老牌主流框架，也會在新模型，新算法，新硬件上表現不佳，所以需要新的框架，在易于擴展同時保持高效?！?/p>

而現在框架呈現出來的能力，的確有超越“前輩”的傾向：

基于Jittor開發的深度學習模型，可以實時自動優化并運行在指定的硬件上，如CPU，GPU，在多種機器視覺任務上能夠比同類產品PyTorch性能提高10%~50%。

團隊還介紹，如此成果，得益于Jittor的兩大創新點：元算子和統一計算圖。這也是Jittor的立身之本。

Jittor的核心：元算子與統一計算圖

根據Jittor官方文檔定義，元算子是指神經網絡所需的基本算子。

在TensorFlow，PyTorch框架底層，有上千個算子，如此多的算子使得開發和優化難度大幅提升。

在設計Jittor的時候，他們就定下了一個目標，即用戶只需要數行代碼，就可定義新的算子和模型。同時在保證易用的同時，不喪失任何可定制性。

所以在Jittor中，多個元算子之間，可以相互融合成更加復雜的算子，這些復雜算子構成了神經網絡計算的多個模塊，如卷積層，歸一化層等等。

他們將這種融合稱為元算子融合，可以提升性能，節省資源。在文檔中，他們分享了一個案例：只用4個元算子，就實現了卷積操作。

他們介紹稱，元算子的可拓展性很強，通過對元算子的簡單修改，就可以實現更多復雜的卷積操作，如擴張卷積、深度卷積、點卷積、分離式卷積、反卷積等。

而且，通過元算子反向傳播閉包，能自動生成反向卷積層。具體如下圖所示，反向卷積層將來自輸出的梯度，通過4個元算子，將梯度反向傳播給卷積層的權重：

Jittor開發團隊介紹稱，在這樣的設計下，元算子和Numpy一樣易于使用，并且超越Numpy能夠實現更復雜更高效的操作。

而且，通過元算子的反向傳播閉包，Jittor可以對所有前向反向算子進行統一管理，這就是他們所說的第二個創新點：統一計算圖。

簡單來說，統一計算圖是完成了多種統一的動態計算圖。根據官方文檔介紹，在Jittor中，核心有四個方面的統一：

• 統一管理前向反向計算圖，使得高階導數可以被支持。
• 統一管理CPU，GPU內存，使得訓練模型時，可以突破原有的GPU顯存限制，讓CPU，GPU可以共享內存。
• 統一同步、異步運行接口，使得數據讀取，內存拷貝，模型計算可以同時進行，提升性能
• 統一管理多次迭代的計算圖，使得平臺可以實現跨迭代的融合優化。

基于這個方面，他們給出了Jittor與其他各個框架的特性對比：

自動微分、動態圖方面，Tensorflow、Pytorch和Jittor都支持。但在同步接口和異步接口方面，Jittor的優異性得到了體現。

同步接口易于編程，異步接口有助于提高性能，Jittor同時支持這兩種接口。

相比之下，Tensorflow部分算子支持統一內存管理，而PyTorch不支持異步接口，而Jittor的所有算子都支持統一內存管理，當深度學習模型將GPU內存資源耗盡時，將使用CPU內存來彌補。

除此之外，Jittor還支持跨迭代融合。

在這些特性的支持下，Jittor具備了動態編譯的能力。

官方文檔介紹稱，通過內置元算子編譯器，可以將用戶用元算子編寫的Python代碼，動態編譯成高性能的C++代碼。

比如，下圖中的Python代碼編寫了神經網絡中常用的批歸一化層（batch norm）， 通過元算子編譯器，動態生成了批歸一化層C++代碼。

開發團隊介紹稱，Jittor還會使用內置的編譯優化，以及LLVM兼容的優化編譯遍（complier pass）來優化動態生成的代碼。

這些編譯會根據硬件設備，自動優化動態編譯的代碼，常見的優化編譯有：循環重排，循環分裂，循環融合，數據打包，向量化，GPU并行。

他們說，這些編譯遍，能夠對C++代碼進一步優化，生成對計算設備友好的底層算子，從而提高性能。

這體現了他們設計Jittor的另一個理念：

所有代碼都是即時編譯并且運行的，包括Jittor本身。用戶可以隨時對Jittor的所有代碼進行修改，并且動態運行。

此外，在整體設計中，他們還遵循了實現與優化分離的理念。

如此打造出來的整體架構，“用戶可以通過前端接口專注于實現，而實現自動被后端優化。從而提升前端代碼的可讀性，以及后端優化的魯棒性和可重用性?！彼麄兘榻B稱。

Jittor的整體架構與上手樣例

具體來說，Jittor的整體架構一共分為四層，如下圖所示：

它是基于Jit編譯技術、完全重新設計的深度學習框架，從上到下分別是應用層，前端層，后端層，硬件層，官方文檔的介紹如下：

應用層的代碼，用戶使用Python編寫，并可以訪問從前端層公開的所有接口。

前端層?是Jittor的組件之一，代碼用Python編寫，提供了元算子的調用接口，來操作Jittor變量和Jittor實現的通用模型。

后端層是Jittor的內核，由C++編寫，管理底層硬件資源。該層包含很多模塊，比如算子融合器、第三方算子、JIT編譯器、統一計算圖、統一內存調度、統一執行器等。

硬件層支持的硬件有CPU和Nvidia GPU。但如果需要讓Jittor支持新的硬件，只需要重載編譯接口即可，讓Jittor移植到新的硬件的難度將大大降低。Jittor開發團隊說，他們將在未來支持更多的計算設備。

如此架構，用起來怎樣？官方文檔介紹稱，從頭只需要若干行代碼，就能訓練一個兩層神經網絡。

上面的代碼，定義了激活函數和全連接層。Jittor開發團隊介紹稱，其實這些層已經集成在了框架中，并使用了類似的實現方式，在這里重新定義，用于更好展示內部機制和實現。

從代碼中可以看出，Jittor的接口和現在主流深度學習框架接口類似，都是使用模塊化的方式定義模型。其中，random、matmul、exp都是Jittor的內置算子。

基于JIT編譯，Jittor的后端會將這幾個算子自動融合成一個算子。

上面的代碼，定義了雙層神經網絡。隱層的神經元個數是10， 使用的激活函數是上面定義好的sigmoid。

最后，可以從頭開始訓練模型。在這段代碼，使用了梯度下降和L2 loss來訓練網絡。訓練過程是異步的。

Jittor開發團隊介紹稱，Jittor會自動計算梯度并且將計算圖保存起來，后端的JIT編譯器會根據計算圖，同時使用算子級別優化和圖級別的優化。

他們進一步解釋稱，在這一示例中，Jittor使用了以下幾種優化：

算子融合：激活函數和loss函數會被融合在一起。

并行化：算子會自動并行化以提升性能和計算密集度，在現代多核CPU和GPU上十分有效。

并發：沒有依賴關系的操作會被并發執行，比如內存拷貝和計算可以并發并相互重疊。

元算子與統一計算圖加持，整體框架優化下，Jittor在一些任務上展現出了性能提升，在多種機器視覺任務上尤為明顯。

多個視覺任務上，性能超過現有主流框架

Jittor開發團隊提供了實驗數據。在ImageNet數據集上，使用Resnet50模型，GPU圖像分類任務性能比PyTorch相比，提升32%；CPU圖像分類任務提升11%。

在CelebA數據集上，使用LSGAN模型，使用GPU處理圖像生成任務，Jittor比PyTorch性能提升達51%。

此外，為了方便更多人上手Jittor，開發團隊采用了和PyTorch較為相似的模塊化接口，并提供輔助轉換腳本，可以將PyTorch的模型自動轉換成Jittor的模型。

他們介紹稱，在參數保存和數據傳輸上，Jittor使用和PyTorch一樣的 Numpy+pickle 協議，所以Jittor和PyTorch的模型可以相互加載和調用。

當然， Jittor作為一個新興深度學習框架，在一些功能上，仍舊需要持續迭代完善。比如生態的建設，以及更大范圍的推廣，仍舊需要很多的努力。

Jittor開發團隊介紹稱，就目前來看，Jittor框架的模型支持還待完善，分布式功能待完善。這也是他們下一階段研發的重點。

首個中國高校深度學習開源框架，清華教授領銜打造

最后，是時候介紹Jittor的開發團隊出場，他們來自清華大學計算機系圖形學實驗室，牽頭者是清華大學計算機系胡事民教授。

該實驗室的主要研究方向是計算機圖形學、計算機視覺、智能信息處理、智能機器人、系統軟件等，在ACM TOG, IEEE TVCG, IEEE PAMI, ACM SIGGRAPH, IEEE CVPR, IEEE ICRA, USENIX ATC等重要國際刊物上發表論文100余篇。

開發Jittor的主力，是該實驗室梁盾、楊國燁、楊國煒、周文洋等一批博士生。

據梁盾透露，他們接下來的計劃，是先圍繞學界，重點發力。希望能成為國內以及世界上學術界最受歡迎，使用最多的框架，并對AI產業界產生積極的影響。

但想要走通這條路，并沒有那么容易。TensorFlow和PyTorch已經成為了當前主流的深度學習框架，正在被全世界的研究者們采用。尤其是PyTorch，正在大面積搶占學術界。

其實從模型特性，以及設計理念來看，PyTorch可能是Jittor更直接的對標對象。

Jittor將如何發力？

在他們的規劃中，接下來將組建開源社區，除了完善框架外，還會聯合多所高校使用Jittor教授課程，以現有人員作為核心，壯大開發團隊和用戶，首要目標是服務更多研究人員。

據說，已經有多位高校教授，決定要在自己課堂上使用。

同時，另一個公開信息也值得關注：胡事民教授從2010年開始，就擔任清華大學—騰訊聯合實驗室主任。在Jittor研發過程中，還得到了這一實驗室支持。

所以這一框架是否會與騰訊展開合作？

目前研究團隊沒有給出直接明確的答復，但表示：非常希望能和更多的產業界的伙伴們聯手推動Jittor的發展。

總之，打造AI框架本身已不易，開源之后更要接受各方直接檢驗。

現在，清華邁出了關鍵一步，虛的不多說，各位收好下方傳送門，走過路過不要錯過，都參與檢驗一下吧~

開源傳送門

https://cg.cs.tsinghua.edu.cn/jittor/

— 完 —

量子位 QbitAI · 頭條號簽約

關注我們，第一時間獲知前沿科技動態