本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

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

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

這一次跟寒武紀(jì)的合作，使Jittor在backbone網(wǎng)絡(luò)模型中的動(dòng)態(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倍。

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

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

還是「元算子」的功勞

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

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

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

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

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

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

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

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

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

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

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

從Jittor元算子到BANG

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

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

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

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

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

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

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

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

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

清華自研，首個(gè)中國(guó)高校深度學(xué)習(xí)開(kāi)源框架

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

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

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

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

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

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

最近，他們還開(kāi)源了一個(gè)智能P圖神器DeepFaceEditing，可以通過(guò)草圖自由編輯人臉。

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

GitHub開(kāi)源：https://github.com/IGLICT/DeepFaceEditing-Jittor