清華開(kāi)源Jittor:首個(gè)國(guó)內(nèi)高校自研深度學(xué)習(xí)框架,一鍵轉(zhuǎn)換PyTorch
繼 Theano、Caffe 之后,又一個(gè)由高校主導(dǎo)的深度學(xué)習(xí)框架開(kāi)源了,而且還是國(guó)產(chǎn)。
深度學(xué)習(xí)框架越來(lái)越多,主導(dǎo)的團(tuán)隊(duì)也從高校研究機(jī)構(gòu)漸漸轉(zhuǎn)向了科技巨頭。但是,學(xué)界在這一領(lǐng)域的力量不容忽視。今日。清華大學(xué)開(kāi)發(fā)了一個(gè)名為計(jì)圖(Jittor)的深度學(xué)習(xí)框架。這一框架有望為深度學(xué)習(xí)社區(qū)提供新的方案,也能夠推動(dòng)深度學(xué)習(xí)框架國(guó)產(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)介紹說(shuō):「研究團(tuán)隊(duì)將神經(jīng)網(wǎng)絡(luò)所需的基本算子定義為元算子」,元算子非常底層,通過(guò)相互融合可以完成復(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)化?;谠阕娱_(kāi)發(fā)的深度學(xué)習(xí)模型,可以被計(jì)圖實(shí)時(shí)地自動(dòng)優(yōu)化并且運(yùn)行在指定的硬件上,如 CPU、GPU。
目前 Jittor 已開(kāi)源,用戶可以采用 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)化分離:可以通過(guò)前端接口專注于實(shí)現(xiàn),而實(shí)現(xiàn)自動(dòng)被后端優(yōu)化。
- 所有都是即時(shí)的:Jittor 的所有代碼都是即時(shí)編譯并且運(yùn)行的,包括 Jittor 本身。用戶可以隨時(shí)對(duì) Jittor 的所有代碼進(jìn)行修改,并且動(dòng)態(tài)運(yùn)行。
Jittor 作為全新的深度學(xué)習(xí)框架,有幾大新特性值得關(guān)注。首先,Jittor 采用了元算子的概念,將各種基本計(jì)算定義為元算子,并通過(guò)結(jié)合不同元算子,實(shí)現(xiàn)深度學(xué)習(xí)中的各項(xiàng)算子功能。
元算子融合。
其次,在算子的設(shè)置上,團(tuán)隊(duì)將元算子的反向傳播進(jìn)行了閉包,即元算子的反向傳播也是元算子。這樣避免了重復(fù)開(kāi)發(fā)。此外,還支持計(jì)算任意高階導(dǎo)數(shù)。
高階導(dǎo)數(shù)及反向傳播閉包。
然后,在編程語(yǔ)言上,Jittor 采用了靈活而易用的 Python。用戶可以使用它,編寫(xiě)元算子計(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)化過(guò)程會(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ì)造成性能損失。從而帶來(lái)了易用性和高效率。
高效同步異步接口。
最后值得注意的是,為了幫助 Jittor 的發(fā)展,并吸收優(yōu)秀框架的特性。Jittor 采用了和 PyTorch 相似的模塊化接口,并有輔助轉(zhuǎn)換腳本,還有和 PyTorch 一樣的 Numpy+pickle 協(xié)議,使得兩者的模型可以互相加載和使用。
模型遷移。
有了如此眾多的特性,估計(jì)用來(lái)寫(xiě)模型會(huì)有非常好的體驗(yàn)。此外,Jittor 使用 Python 和 C ++編寫(xiě),也支持 GPU 與 CUDA。
測(cè)評(píng)結(jié)果
Jittor 的性能如何,官方也提供了測(cè)評(píng)結(jié)果進(jìn)行參考。目前 ResNet、VGG、SSD、DeepLab、LSGAN 等多個(gè)網(wǎng)絡(luò)模型已經(jīng)在 Jittor 平臺(tái)實(shí)現(xiàn),可供用戶使用。與同類型框架相比,Jittor 在收斂精度一致情況下,推理速度取得了 10%-50% 的性能提升。
示例代碼
Jittor 前端語(yǔ)言為 Python。前端使用模塊化的設(shè)計(jì),類似于 PyTorch,Keras,后端則使用高性能語(yǔ)言編寫(xiě),如 CUDA,C++。Jittor 官網(wǎng)目前已提供了示例代碼,我們可以來(lái)分析一下。
下面的代碼演示了如何一步一步使用 Python 代碼,從頭對(duì)一個(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è)簡(jiǎn)單的全連接神經(jīng)網(wǎng)絡(luò)的代碼。可以看到,類似于 PyTorch,只需要導(dǎo)入包,采用類繼承的方式定義模型、數(shù)據(jù)處理(分批等)方式,并設(shè)置訓(xùn)練循環(huán)即可。從上手難度來(lái)看,熟悉 PyTorch 的用戶可以直接上手,甚至沒(méi)有因?yàn)樽兞棵煌鴰?lái)影響。
研究團(tuán)隊(duì)
據(jù)官網(wǎng)介紹,Jittor 的開(kāi)發(fā)團(tuán)隊(duì)來(lái)自清華大學(xué)計(jì)算機(jī)系的圖形學(xué)實(shí)驗(yàn)室,實(shí)驗(yàn)室負(fù)責(zé)人為胡事民教授。
該實(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ù)國(guó)家研究中心下的可視媒體智能計(jì)算團(tuán)隊(duì)。目前有教授 2 名、副教授 3 名、助理研究員 1 名、博士后 4 名和研究生 50 多名。
Jittor 項(xiàng)目的主要研發(fā)主力為實(shí)驗(yàn)室的梁盾、楊國(guó)燁、楊國(guó)煒和周文洋等等博士生,此外該項(xiàng)目也得到了清華-騰訊聯(lián)合實(shí)驗(yàn)室的資助和支持。