清華「計(jì)圖」現(xiàn)在支持國產(chǎn)芯片了！動態(tài)圖推理比PyTorch快了270倍

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

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

清華自研的深度學(xué)習(xí)框架計(jì)圖（Jittor）在動態(tài)圖推理速度上又一次完勝PyTorch。

最近，計(jì)圖團(tuán)隊(duì)完成了在寒武紀(jì)芯片MLU270上的移植。

這一次跟寒武紀(jì)的合作，使Jittor在backbone網(wǎng)絡(luò)模型中的動態(tài)圖推理速度較PyTorch平均提升了276.69倍。

從團(tuán)隊(duì)公布的實(shí)驗(yàn)結(jié)果可以看到，在寒武紀(jì)芯片上分別用計(jì)圖（Jittor）和PyTorch進(jìn)行推理。

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

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

平均提升速度達(dá)276.69倍。

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

團(tuán)隊(duì)以resnet18網(wǎng)絡(luò)為例進(jìn)行對比，baseline是CPU；結(jié)果顯示：計(jì)圖的精度損失比PyTorch降低了42.53%。

還是「元算子」的功勞

性能得到如此大的提升，是怎么實(shí)現(xiàn)的呢？

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

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

根據(jù)Jittor官方文檔定義，元算子是指神經(jīng)網(wǎng)絡(luò)所需的基本算子。

早在設(shè)計(jì)Jittor的時候，研發(fā)團(tuán)隊(duì)就定下了一個目標(biāo)，即用戶只需要數(shù)行代碼，就可定義新的算子和模型。在保證易用的同時，不喪失任何可定制性。

所以在Jittor中，多個元算子之間，可以相互融合成更加復(fù)雜的算子，這些復(fù)雜算子構(gòu)成了神經(jīng)網(wǎng)絡(luò)計(jì)算的多個模塊，如卷積層，歸一化層等等。

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

研究團(tuán)隊(duì)把神經(jīng)網(wǎng)絡(luò)所需的基本算子，定義為三類共18個元算子。

讓元算子相互融合形成常用算子，這樣就能通過優(yōu)化非常少的元算子，實(shí)現(xiàn)對常用算子性能的提升。

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

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

從Jittor元算子到BANG

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

BANG語言是專門針對寒武紀(jì)產(chǎn)品架構(gòu)的一種代碼語言，它能極大優(yōu)化寒武紀(jì)芯片的通用編程能力，提升用戶編程的自由度。

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

Jittor內(nèi)置的元算子編譯器可以把元算子自動轉(zhuǎn)化為BANG算子。

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

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

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

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

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

清華自研，首個中國高校深度學(xué)習(xí)開源框架

計(jì)圖（Jittor）的開發(fā)團(tuán)隊(duì)，均來自清華大學(xué)計(jì)算機(jī)系圖形學(xué)實(shí)驗(yàn)室，負(fù)責(zé)人是清華大學(xué)計(jì)算機(jī)系的胡事民教授。

而主要負(fù)責(zé)開發(fā)的，則是來自實(shí)驗(yàn)室的博士生們：梁盾、楊國燁、楊國煒、周文洋……

計(jì)圖（Jittor）與主流的深度學(xué)習(xí)框架TensorFlow、Pytorch等最大的不同在于，它是一個完全基于動態(tài)編譯（Just-in-time）、使用元算子和統(tǒng)一計(jì)算圖的深度學(xué)習(xí)框架。

它可以不像Pytorch那樣依賴tracing機(jī)制，讓用戶的操作更加簡單，同時還可以得到更好的學(xué)習(xí)訓(xùn)練效果。

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

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

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

感興趣的同學(xué)可以去試試哦~

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

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