清華開源Jittor：首個國內高校自研深度學習框架，一鍵轉換PyTorch
繼 Theano、Caffe 之后，又一個由高校主導的深度學習框架開源了，而且還是國產。
深度學習框架越來越多，主導的團隊也從高校研究機構漸漸轉向了科技巨頭。但是，學界在這一領域的力量不容忽視。今日。清華大學開發(fā)了一個名為計圖（Jittor）的深度學習框架。這一框架有望為深度學習社區(qū)提供新的方案，也能夠推動深度學習框架國產化的進程。
據悉，計圖（Jittor：Just in Time）是一個采用元算子表達神經網絡計算單元、完全基于動態(tài)編譯（Just-in-Time）的深度學習框架，其主要特性為元算子和統(tǒng)一計算圖。
據官網介紹說：「研究團隊將神經網絡所需的基本算子定義為元算子」，元算子非常底層，通過相互融合可以完成復雜的深度學習計算，體現(xiàn)出了易于使用的特點。并且研發(fā)團隊表示，采用元算子的 Jittor 已超越 Numpy，能夠實現(xiàn)更復雜更高效的操作。
而另一層面，統(tǒng)一計算圖則是融合了靜態(tài)計算圖和動態(tài)計算圖的諸多優(yōu)點，在易于使用的同時，提供高性能的優(yōu)化?；谠阕娱_發(fā)的深度學習模型，可以被計圖實時地自動優(yōu)化并且運行在指定的硬件上，如 CPU、GPU。
目前 Jittor 已開源，用戶可以采用 pip 等方法下載使用。
Jittor 官網：https://cg.cs.tsinghua.edu.cn/jittor/
項目地址：https://github.com/Jittor/jittor
即時、易用、可定制：實現(xiàn)和優(yōu)化分離的框架
據官網介紹，Jittor 在設計時秉持易用、靈活、即時的設計理念：
- 易用且可定制：只需要數行代碼，就可定義新的算子和模型。
- 實現(xiàn)與優(yōu)化分離：可以通過前端接口專注于實現(xiàn)，而實現(xiàn)自動被后端優(yōu)化。
- 所有都是即時的：Jittor 的所有代碼都是即時編譯并且運行的，包括 Jittor 本身。用戶可以隨時對 Jittor 的所有代碼進行修改，并且動態(tài)運行。
Jittor 作為全新的深度學習框架，有幾大新特性值得關注。首先，Jittor 采用了元算子的概念，將各種基本計算定義為元算子，并通過結合不同元算子，實現(xiàn)深度學習中的各項算子功能。
元算子融合。
其次，在算子的設置上，團隊將元算子的反向傳播進行了閉包，即元算子的反向傳播也是元算子。這樣避免了重復開發(fā)。此外，還支持計算任意高階導數。
高階導數及反向傳播閉包。
然后，在編程語言上，Jittor 采用了靈活而易用的 Python。用戶可以使用它，編寫元算子計算的 Python 代碼，然后 Jittor
將其動態(tài)編譯為 C++，實現(xiàn)高性能。
算子動態(tài)編譯。
更重要的是，Jittor 內置了優(yōu)化編譯遍（complier pass）功能，和 LLVM 兼容。優(yōu)化過程會根據運行的硬件而定，以下為已支持的優(yōu)化編譯遍。
自動優(yōu)化。
在內存管理上，CPU 和 GPU 內存統(tǒng)一管理。GPU 內存不夠的時候會調用 GPU。
統(tǒng)一內存管理。
而在接口方面，Jittor 同時提供了同步和異步接口，切換不會造成性能損失。從而帶來了易用性和高效率。
高效同步異步接口。
最后值得注意的是，為了幫助 Jittor 的發(fā)展，并吸收優(yōu)秀框架的特性。Jittor 采用了和 PyTorch 相似的模塊化接口，并有輔助轉換腳本，還有和 PyTorch 一樣的 Numpy+pickle 協(xié)議，使得兩者的模型可以互相加載和使用。
模型遷移。
有了如此眾多的特性，估計用來寫模型會有非常好的體驗。此外，Jittor 使用 Python 和 C ++編寫，也支持 GPU 與 CUDA。
測評結果
Jittor 的性能如何，官方也提供了測評結果進行參考。目前 ResNet、VGG、SSD、DeepLab、LSGAN 等多個網絡模型已經在 Jittor 平臺實現(xiàn)，可供用戶使用。與同類型框架相比，Jittor 在收斂精度一致情況下，推理速度取得了 10%-50% 的性能提升。
示例代碼
Jittor 前端語言為 Python。前端使用模塊化的設計，類似于 PyTorch，Keras，后端則使用高性能語言編寫，如 CUDA，C++。Jittor 官網目前已提供了示例代碼，我們可以來分析一下。
下面的代碼演示了如何一步一步使用 Python 代碼，從頭對一個雙層神經網絡建模。
- 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()}")
以上為實現(xiàn)一個簡單的全連接神經網絡的代碼?？梢钥吹剑愃朴?PyTorch，只需要導入包，采用類繼承的方式定義模型、數據處理（分批等）方式，并設置訓練循環(huán)即可。從上手難度來看，熟悉 PyTorch 的用戶可以直接上手，甚至沒有因為變量名不同而帶來影響。
研究團隊
據官網介紹，Jittor 的開發(fā)團隊來自清華大學計算機系的圖形學實驗室，實驗室負責人為胡事民教授。
該實驗室成立于 1998 年 3 月，2007 年發(fā)展成為清華大學可視媒體研究中心，2010 年獲批成為北京市工程技術研究中心，同年和騰訊公司合作成立清華-騰訊互聯(lián)網創(chuàng)新技術聯(lián)合實驗室，并于 2018 年，成立北京信息科學與技術國家研究中心下的可視媒體智能計算團隊。目前有教授 2 名、副教授 3 名、助理研究員 1 名、博士后 4 名和研究生 50 多名。
Jittor 項目的主要研發(fā)主力為實驗室的梁盾、楊國燁、楊國煒和周文洋等等博士生，此外該項目也得到了清華-騰訊聯(lián)合實驗室的資助和支持。
