開源圖神經(jīng)網(wǎng)絡(luò)框架DGL升級(jí):GCMC訓(xùn)練時(shí)間從1天縮到1小時(shí)
本文經(jīng)AI新媒體量子位(公眾號(hào)ID:QbitAI)授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請(qǐng)聯(lián)系出處。
又一個(gè)AI框架迎來升級(jí)。
這次,是紐約大學(xué)、亞馬遜聯(lián)手推出圖神經(jīng)網(wǎng)絡(luò)框架DGL。

不僅全面上線了對(duì)異構(gòu)圖的支持,復(fù)現(xiàn)并開源了相關(guān)異構(gòu)圖神經(jīng)網(wǎng)絡(luò)的代碼,在GCMC、RCGN等業(yè)內(nèi)知名的模型實(shí)現(xiàn)上,取得更好效果。
還發(fā)布了訓(xùn)練知識(shí)圖譜嵌入(Knowledge Graph Embedding)專用包DGL-KE,在許多經(jīng)典的圖嵌入模型上都獲得了性能提升。
目前,這個(gè)在2018年NeurIPS大會(huì)上發(fā)布的框架,已經(jīng)獲得了3000+標(biāo)星,已經(jīng)是圖神經(jīng)網(wǎng)絡(luò)和圖機(jī)器學(xué)習(xí)領(lǐng)域較受關(guān)注的工具。
異構(gòu)圖:神經(jīng)網(wǎng)絡(luò)的新方向
異構(gòu)圖是一個(gè)與同構(gòu)圖相對(duì)應(yīng)的新概念。
傳統(tǒng)同構(gòu)圖(Homogeneous Graph)數(shù)據(jù)中只存在一種節(jié)點(diǎn)和邊,因此在構(gòu)建圖神經(jīng)網(wǎng)絡(luò)時(shí)所有節(jié)點(diǎn)共享同樣的模型參數(shù)并且擁有同樣維度的特征空間。
而異構(gòu)圖(Heterogeneous Graph)中可以存在不只一種節(jié)點(diǎn)和邊,因此允許不同類型的節(jié)點(diǎn)擁有不同維度的特征或?qū)傩浴?/p>
這一特點(diǎn)使得異構(gòu)圖的應(yīng)用十分廣泛,如果用圖來描述我們和周圍事物的關(guān)系就會(huì)發(fā)現(xiàn)所產(chǎn)生的圖都是天然異構(gòu)的,比如這個(gè)例子:
我今天看了電影《流浪地球》,那“我”作為觀眾和電影《流浪地球》之間就建立了“看了”這一關(guān)系。
異構(gòu)圖可以用來描述這種交互關(guān)系的集合。這個(gè)圖分“觀眾”和“電影”兩類節(jié)點(diǎn),以及“看了”這一類邊。
“我”作為觀眾,和電影所具有的屬性一定是不同的,需要用不同的模型或者不同的特征維度來表達(dá)。
所以說,這張圖就天然具有了異構(gòu)性。而且,基于異構(gòu)圖訓(xùn)練的神經(jīng)網(wǎng)絡(luò)與傳統(tǒng)方法相比,在一些場(chǎng)景中,也能夠獲得更好的效果與表現(xiàn)。
現(xiàn)在已經(jīng)被應(yīng)用到了知識(shí)圖譜、推薦系統(tǒng)以及惡意賬戶識(shí)別等領(lǐng)域以及任務(wù)中。

其中比較有代表性的兩個(gè)模型分別是用于節(jié)點(diǎn)分類和鏈接預(yù)測(cè)等任務(wù)的RGCN和用于產(chǎn)品推薦的GCMC。
但如何設(shè)計(jì)“既快又好”的深度神經(jīng)網(wǎng)絡(luò)?
這正是紐約大學(xué)、亞馬遜聯(lián)手推出圖神經(jīng)網(wǎng)絡(luò)框架DGL的出發(fā)點(diǎn)。
新版DGL:異構(gòu)圖神經(jīng)網(wǎng)絡(luò)利器
這次更新放出的,是DGL 0.4版本,讓整個(gè)框架變得更加實(shí)用。
不僅全面上線對(duì)于異構(gòu)圖的支持,還復(fù)現(xiàn)并開源了相關(guān)異構(gòu)圖神經(jīng)網(wǎng)絡(luò)的代碼:

△DGL 0.4中基于異構(gòu)圖的RGCN層實(shí)現(xiàn)代碼
官方表示,在這一新版本的DGL上實(shí)現(xiàn)業(yè)內(nèi)比較突出的異構(gòu)圖神經(jīng)網(wǎng)絡(luò),性能上也有更好的表現(xiàn):
GCMC:DGL的實(shí)現(xiàn)相比原作者實(shí)現(xiàn)在MovieLens-100K上有5倍加速,在MovieLens-1M上有22倍加速。DGL的內(nèi)存優(yōu)化支持在一塊GPU上對(duì)MovieLens-10M進(jìn)行訓(xùn)練(原實(shí)現(xiàn)需要從CPU動(dòng)態(tài)加載數(shù)據(jù)),從而將原本需要24小時(shí)的訓(xùn)練時(shí)間縮短到了1個(gè)多小時(shí)。
RGCN:使用全新的異構(gòu)圖接口重新實(shí)現(xiàn)了RGCN。新的實(shí)現(xiàn)大大減少了內(nèi)存開銷。原實(shí)現(xiàn)在AM數(shù)據(jù)集上(邊數(shù)>5M)由于內(nèi)存開銷太大,只能在CPU上計(jì)算,而DGL能使用GPU進(jìn)行加速并獲得了291倍的速度提升。
HAN:提供了靈活的接口可以將一個(gè)異構(gòu)圖通過元路徑(metapath)轉(zhuǎn)變成同構(gòu)圖。
Metapath2vec:新的元路徑采樣實(shí)現(xiàn)比原作者的實(shí)現(xiàn)快兩倍。
而且, DGL也發(fā)布了針對(duì)分子化學(xué)的模型庫DGL-Chem,以及訓(xùn)練知識(shí)圖譜嵌入(Knowledge Graph Embedding)專用包DGL-KE。DGL-Chem提供了包括分子性質(zhì)預(yù)測(cè)和分子結(jié)構(gòu)生成等預(yù)訓(xùn)練模型。
在單GPU上,DGL-KE能在7分鐘內(nèi)使用經(jīng)典的TransE模型訓(xùn)練出FB15K的圖嵌入。相比GraphVite(v0.1.0)在4塊GPU上需要14分鐘。
DGL-KE首個(gè)版本發(fā)布了TransE、ComplEx和Distmult模型,支持CPU訓(xùn)練、GPU訓(xùn)練、CPU和GPU混合訓(xùn)練以及單機(jī)多進(jìn)程訓(xùn)練。
DGL團(tuán)隊(duì)表示,在未來也會(huì)發(fā)布更多模型和訓(xùn)練方式。