乾明 發自 凹非寺
量子位 報道 | 公眾號 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/
— 完 —
量子位 QbitAI · 頭條號簽約
關注我們，第一時間獲知前沿科技動態
閱讀更多 量子位 的文章
