繼 Theano、Caffe 之后，又一個(gè)由高校主導(dǎo)的深度學(xué)習(xí)框架開源了，而且還是國產(chǎn)。

深度學(xué)習(xí)框架越來越多，主導(dǎo)的團(tuán)隊(duì)也從高校研究機(jī)構(gòu)漸漸轉(zhuǎn)向了科技巨頭。但是，學(xué)界在這一領(lǐng)域的力量不容忽視。今日。清華大學(xué)開發(fā)了一個(gè)名為計(jì)圖（Jittor）的深度學(xué)習(xí)框架。這一框架有望為深度學(xué)習(xí)社區(qū)提供新的方案，也能夠推動(dòng)深度學(xué)習(xí)框架國產(chǎn)化的進(jìn)程。

據(jù)悉，計(jì)圖（Jittor：Just in Time）是一個(gè)采用元算子表達(dá)神經(jīng)網(wǎng)絡(luò)計(jì)算單元、完全基于動(dòng)態(tài)編譯（Just-in-Time）的深度學(xué)習(xí)框架，其主要特性為元算子和統(tǒng)一計(jì)算圖。

據(jù)官網(wǎng)介紹說：「研究團(tuán)隊(duì)將神經(jīng)網(wǎng)絡(luò)所需的基本算子定義為元算子」，元算子非常底層，通過相互融合可以完成復(fù)雜的深度學(xué)習(xí)計(jì)算，體現(xiàn)出了易于使用的特點(diǎn)。并且研發(fā)團(tuán)隊(duì)表示，采用元算子的 Jittor 已超越 Numpy，能夠?qū)崿F(xiàn)更復(fù)雜更高效的操作。

而另一層面，統(tǒng)一計(jì)算圖則是融合了靜態(tài)計(jì)算圖和動(dòng)態(tài)計(jì)算圖的諸多優(yōu)點(diǎn)，在易于使用的同時(shí)，提供高性能的優(yōu)化。基于元算子開發(fā)的深度學(xué)習(xí)模型，可以被計(jì)圖實(shí)時(shí)地自動(dòng)優(yōu)化并且運(yùn)行在指定的硬件上，如 CPU、GPU。

目前 Jittor 已開源，用戶可以采用 pip 等方法下載使用。

Jittor 官網(wǎng)：https://cg.cs.tsinghua.edu.cn/jittor/

項(xiàng)目地址：https://github.com/Jittor/jittor

即時(shí)、易用、可定制：實(shí)現(xiàn)和優(yōu)化分離的框架

據(jù)官網(wǎng)介紹，Jittor 在設(shè)計(jì)時(shí)秉持易用、靈活、即時(shí)的設(shè)計(jì)理念：

• 易用且可定制：只需要數(shù)行代碼，就可定義新的算子和模型。
• 實(shí)現(xiàn)與優(yōu)化分離：可以通過前端接口專注于實(shí)現(xiàn)，而實(shí)現(xiàn)自動(dòng)被后端優(yōu)化。
• 所有都是即時(shí)的：Jittor 的所有代碼都是即時(shí)編譯并且運(yùn)行的，包括 Jittor 本身。用戶可以隨時(shí)對 Jittor 的所有代碼進(jìn)行修改，并且動(dòng)態(tài)運(yùn)行。

Jittor 作為全新的深度學(xué)習(xí)框架，有幾大新特性值得關(guān)注。首先，Jittor 采用了元算子的概念，將各種基本計(jì)算定義為元算子，并通過結(jié)合不同元算子，實(shí)現(xiàn)深度學(xué)習(xí)中的各項(xiàng)算子功能。

元算子融合。

其次，在算子的設(shè)置上，團(tuán)隊(duì)將元算子的反向傳播進(jìn)行了閉包，即元算子的反向傳播也是元算子。這樣避免了重復(fù)開發(fā)。此外，還支持計(jì)算任意高階導(dǎo)數(shù)。

高階導(dǎo)數(shù)及反向傳播閉包。

然后，在編程語言上，Jittor 采用了靈活而易用的 Python。用戶可以使用它，編寫元算子計(jì)算的 Python 代碼，然后 Jittor

將其動(dòng)態(tài)編譯為 C++，實(shí)現(xiàn)高性能。

算子動(dòng)態(tài)編譯。

更重要的是，Jittor 內(nèi)置了優(yōu)化編譯遍（complier pass）功能，和 LLVM 兼容。優(yōu)化過程會(huì)根據(jù)運(yùn)行的硬件而定，以下為已支持的優(yōu)化編譯遍。

自動(dòng)優(yōu)化。

在內(nèi)存管理上，CPU 和 GPU 內(nèi)存統(tǒng)一管理。GPU 內(nèi)存不夠的時(shí)候會(huì)調(diào)用 GPU。

統(tǒng)一內(nèi)存管理。

而在接口方面，Jittor 同時(shí)提供了同步和異步接口，切換不會(huì)造成性能損失。從而帶來了易用性和高效率。

高效同步異步接口。

最后值得注意的是，為了幫助 Jittor 的發(fā)展，并吸收優(yōu)秀框架的特性。Jittor 采用了和 PyTorch 相似的模塊化接口，并有輔助轉(zhuǎn)換腳本，還有和 PyTorch 一樣的 Numpy+pickle 協(xié)議，使得兩者的模型可以互相加載和使用。

模型遷移。

有了如此眾多的特性，估計(jì)用來寫模型會(huì)有非常好的體驗(yàn)。此外，Jittor 使用 Python 和 C ++編寫，也支持 GPU 與 CUDA。

測評結(jié)果

Jittor 的性能如何，官方也提供了測評結(jié)果進(jìn)行參考。目前 ResNet、VGG、SSD、DeepLab、LSGAN 等多個(gè)網(wǎng)絡(luò)模型已經(jīng)在 Jittor 平臺(tái)實(shí)現(xiàn)，可供用戶使用。與同類型框架相比，Jittor 在收斂精度一致情況下，推理速度取得了 10%-50% 的性能提升。

示例代碼

Jittor 前端語言為 Python。前端使用模塊化的設(shè)計(jì)，類似于 PyTorch，Keras，后端則使用高性能語言編寫，如 CUDA，C++。Jittor 官網(wǎng)目前已提供了示例代碼，我們可以來分析一下。

下面的代碼演示了如何一步一步使用 Python 代碼，從頭對一個(gè)雙層神經(jīng)網(wǎng)絡(luò)建模。

import jittor as jt 
 
from jittor import Module 
 
from jittor import nn 
 
class Model(Module): 
 
   def __init__(self): 
 
       self.layer1 = nn.Linear(1, 10) 
 
       self.relu = nn.Relu() 
 
       self.layer2 = nn.Linear(10, 1) 
 
   def execute (self,x) : 
 
       x = self.layer1(x) 
 
       x = self.relu(x) 
 
       x = self.layer2(x) 
 
       return x 
 
def get_data(n): # generate random data for training test. 
 
   for i in range(n): 
 
       x = np.random.rand(batch_size, 1) 
 
       y = x*x 
 
       yield jt.float32(x), jt.float32(y) 
 
model = Model() 
 
learning_rate = 0.1 
 
optim = nn.SGD(model.parameters(), learning_rate) 
 
for i,(x,y) in enumerate(get_data(n)): 
 
   pred_y = model(x) 
 
   loss = ((pred_y - y)**2) 
 
   loss_mean = loss.mean() 
 
   optim.step (loss_mean) 
 
   print(f"step {i}, loss = {loss_mean.data.sum()}") 

以上為實(shí)現(xiàn)一個(gè)簡單的全連接神經(jīng)網(wǎng)絡(luò)的代碼。可以看到，類似于 PyTorch，只需要導(dǎo)入包，采用類繼承的方式定義模型、數(shù)據(jù)處理（分批等）方式，并設(shè)置訓(xùn)練循環(huán)即可。從上手難度來看，熟悉 PyTorch 的用戶可以直接上手，甚至沒有因?yàn)樽兞棵煌鴰碛绊憽?/p>

研究團(tuán)隊(duì)

據(jù)官網(wǎng)介紹，Jittor 的開發(fā)團(tuán)隊(duì)來自清華大學(xué)計(jì)算機(jī)系的圖形學(xué)實(shí)驗(yàn)室，實(shí)驗(yàn)室負(fù)責(zé)人為胡事民教授。

[[319341]]

該實(shí)驗(yàn)室成立于 1998 年 3 月，2007 年發(fā)展成為清華大學(xué)可視媒體研究中心，2010 年獲批成為北京市工程技術(shù)研究中心，同年和騰訊公司合作成立清華-騰訊互聯(lián)網(wǎng)創(chuàng)新技術(shù)聯(lián)合實(shí)驗(yàn)室，并于 2018 年，成立北京信息科學(xué)與技術(shù)國家研究中心下的可視媒體智能計(jì)算團(tuán)隊(duì)。目前有教授 2 名、副教授 3 名、助理研究員 1 名、博士后 4 名和研究生 50 多名。

Jittor 項(xiàng)目的主要研發(fā)主力為實(shí)驗(yàn)室的梁盾、楊國燁、楊國煒和周文洋等等博士生，此外該項(xiàng)目也得到了清華-騰訊聯(lián)合實(shí)驗(yàn)室的資助和支持。