清華「計圖」現在支持國產芯片了！動態圖推理比PyTorch快了270倍

明敏 2021-05-11 15:36:23 來源：量子位

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

清華自研的深度學習框架計圖（Jittor）在動態圖推理速度上又一次完勝PyTorch。

最近，計圖團隊完成了在寒武紀芯片MLU270上的移植。

這一次跟寒武紀的合作，使Jittor在backbone網絡模型中的動態圖推理速度較PyTorch平均提升了276.69倍。

從團隊公布的實驗結果可以看到，在寒武紀芯片上分別用計圖（Jittor）和PyTorch進行推理。

計圖（Jittor）在16種backbone網絡模型中的推理速度都較PyTorch大幅提升，其中包括alexnet、vgg系列、resnet系列。

其中最快的是alexnet，提升速度達到了464.43倍，最慢的resnet50也達到了153倍。

平均提升速度達276.69倍。

與此同時，計圖的精度損失也較PyTorch大幅減少。

團隊以resnet18網絡為例進行對比，baseline是CPU；結果顯示：計圖的精度損失比PyTorch降低了42.53%。

還是「元算子」的功勞

性能得到如此大的提升，是怎么實現的呢？

要知道，在TensorFlow，PyTorch框架底層，有上千個算子；想要移植新的硬件，就必須將大量的算子復寫，這樣一來不僅工作量劇增、而且難度也會加大。

Jittor這次還是依靠他們的核心武器——元算子。

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

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

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

這一次將Jittor移植到寒武紀上，也是同樣的原理。

研究團隊把神經網絡所需的基本算子，定義為三類共18個元算子。

讓元算子相互融合形成常用算子，這樣就能通過優化非常少的元算子，實現對常用算子性能的提升。

并且，元算子還是反向傳播閉包，所有元算子反向傳播后依舊是元算子。

所以在完成三類元算子的移植后，Jittor天然就能支持大部分常用元子的推理和簡單訓練。

從Jittor元算子到BANG

為了能讓用戶更為簡單地操作，Jittor內置了元算子編譯器，可以將用戶的Python代碼動態編譯成寒武紀BANG語言。

BANG語言是專門針對寒武紀產品架構的一種代碼語言，它能極大優化寒武紀芯片的通用編程能力，提升用戶編程的自由度。

并且它還有全套編譯工具鏈來提高性能。包括CNCC（Cambricon Neuware Compiler Collection ）、CNAS（Cambricon Neuware Assembler）、CNLINKER（Cambricon Neuware Linker）。

Jittor內置的元算子編譯器可以把元算子自動轉化為BANG算子。

采用動態編譯的方式，能在運行時獲取更多的信息，比如計算圖上下文、形狀信息等，這些信息都可以進一步提升算子的性能。

并且，Jittor還內置了可以根據硬件設備自動對BANG語言進行優化的優化編譯遍(complier pass)，從而生成對硬件設備更為適合的底層算子。

用這種方法，僅僅一行Python代碼，就能表示BatchNorm算子的核心思想。

元算子編譯器再把這行代碼自動優化成BANG語言代碼。

由于BANG語言的設計更加成熟，提供了類似于CUDA語言的線程調度模式，使得上手和調試都更加簡單；也能更好地釋放寒武紀芯片的算力。

清華自研，首個中國高校深度學習開源框架

計圖（Jittor）的開發團隊，均來自清華大學計算機系圖形學實驗室，負責人是清華大學計算機系的胡事民教授。

而主要負責開發的，則是來自實驗室的博士生們：梁盾、楊國燁、楊國煒、周文洋……

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

它可以不像Pytorch那樣依賴tracing機制，讓用戶的操作更加簡單，同時還可以得到更好的學習訓練效果。

對于這一次取得的成果，開發者之一梁盾透露未來有可能發表在頂會上，并且會開源給大家。

此前，計圖開源的點云模型庫，在多種主流模型上訓練性能較Pytorch提升一倍以上；并發布了第一個支持金屬度、粗糙度的可微渲染庫。

最近，他們還開源了一個智能P圖神器DeepFaceEditing，可以通過草圖自由編輯人臉。

感興趣的同學可以去試試哦~

GitHub開源：https://github.com/IGLICT/DeepFaceEditing-Jittor

參考鏈接：
https://mp.weixin.qq.com/s/LqF2wqTAiLdVKcZ4facxug