原創 關注前沿科技 量子位
乾明 發自 凹非寺
量子位 報導 | 公眾號 QbitAI
AI框架，又來重磅中國玩家。
剛剛，清華自研的深度學習框架，正式對外開源。「貴系」計算機系的圖形實驗室出品，取名Jittor，中文名計圖。
值得一提的是，這也是首個來自中國高校科研機構的開源深度學習框架，之前，業內來自「高校」的還有加拿大蒙特婁大學的Theano，UC伯克利的Caffe。
與主流的深度學習框架TensorFlow、Pytorch不同，Jittor是一個完全基於動態編譯（Just-in-time）、使用元算子和統一計算圖的深度學習框架。
研發團隊介紹稱，開發Jittor是為了將新技術、硬體和模型的能力，更好地釋放出來。
「深度學習發展迅猛，TensorFlow、PyTorch這些老牌主流框架，也會在新模型，新算法，新硬體上表現不佳，所以需要新的框架，在易於擴展同時保持高效。」
而現在框架呈現出來的能力，的確有超越「前輩」的傾向：
基於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的所有代碼進行修改，並且動態運行。
此外，在整體設計中，他們還遵循了實現與優化分離的理念。
如此打造出來的整體架構，「用戶可以通過前端接口專注於實現，而實現自動被後端優化。從而提升前端代碼的可讀性，以及後端優化的魯棒性和可重用性。」他們介紹稱。
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/
作者系網易新聞·網易號「各有態度」籤約作者
— 完 —
北京時間3月21日早10:30，兩位輕舟智航頂級技術專家、前Waymo核心工程師將交流討論現有無人車技術思路，提出無人車創新型技術思路及大規模智能仿真系統的具體應用及實踐。
戳二維碼，備註「無人車」即可報名、加交流群與同好討論交流無人車進展~
高能直播 | 無人車創新型技術思路及仿真系統應用講解
量子位 QbitAI · 頭條號籤約作者
վ'ᴗ' ի 追蹤AI技術和產品新動態
喜歡就點「在看」吧 !
原標題：《清華自研深度學習框架「計圖」開源！多項任務性能超過PyTorch》
閱讀原文