開源圖神經網絡框架DGL升級：GCMC訓練時間從1天縮到1小時，RGCN實現速度提升291倍

乾明 2019-10-15 12:58:28 來源：量子位

乾明 編輯整理

量子位 報道 | 公眾號 QbitAI

又一個AI框架迎來升級。

這次，是紐約大學、亞馬遜聯手推出圖神經網絡框架DGL。

不僅全面上線了對異構圖的支持，復現并開源了相關異構圖神經網絡的代碼，在GCMC、RCGN等業內知名的模型實現上，取得更好效果。

還發布了訓練知識圖譜嵌入（Knowledge Graph Embedding）專用包DGL-KE，在許多經典的圖嵌入模型上都獲得了性能提升。

目前，這個在2018年NeurIPS大會上發布的框架，已經獲得了3000+標星，已經是圖神經網絡和圖機器學習領域較受關注的工具。

異構圖：神經網絡的新方向

異構圖是一個與同構圖相對應的新概念。

傳統同構圖（Homogeneous Graph）數據中只存在一種節點和邊，因此在構建圖神經網絡時所有節點共享同樣的模型參數并且擁有同樣維度的特征空間。

而異構圖（Heterogeneous Graph）中可以存在不只一種節點和邊，因此允許不同類型的節點擁有不同維度的特征或屬性。

這一特點使得異構圖的應用十分廣泛，如果用圖來描述我們和周圍事物的關系就會發現所產生的圖都是天然異構的，比如這個例子：

我今天看了電影《流浪地球》，那“我”作為觀眾和電影《流浪地球》之間就建立了“看了”這一關系。

異構圖可以用來描述這種交互關系的集合。這個圖分“觀眾”和“電影”兩類節點，以及“看了”這一類邊。

“我”作為觀眾，和電影所具有的屬性一定是不同的，需要用不同的模型或者不同的特征維度來表達。

所以說，這張圖就天然具有了異構性。而且，基于異構圖訓練的神經網絡與傳統方法相比，在一些場景中，也能夠獲得更好的效果與表現。

現在已經被應用到了知識圖譜、推薦系統以及惡意賬戶識別等領域以及任務中。

其中比較有代表性的兩個模型分別是用于節點分類和鏈接預測等任務的RGCN和用于產品推薦的GCMC。

但如何設計“既快又好”的深度神經網絡？

這正是紐約大學、亞馬遜聯手推出圖神經網絡框架DGL的出發點。

新版DGL：異構圖神經網絡利器

這次更新放出的，是DGL 0.4版本，讓整個框架變得更加實用。

不僅全面上線對于異構圖的支持，還復現并開源了相關異構圖神經網絡的代碼：

△DGL 0.4中基于異構圖的RGCN層實現代碼

官方表示，在這一新版本的DGL上實現業內比較突出的異構圖神經網絡，性能上也有更好的表現：

GCMC：DGL的實現相比原作者實現在MovieLens-100K上有5倍加速，在MovieLens-1M上有22倍加速。DGL的內存優化支持在一塊GPU上對MovieLens-10M進行訓練（原實現需要從CPU動態加載數據），從而將原本需要24小時的訓練時間縮短到了1個多小時。

RGCN：使用全新的異構圖接口重新實現了RGCN。新的實現大大減少了內存開銷。原實現在AM數據集上（邊數>5M）由于內存開銷太大，只能在CPU上計算，而DGL能使用GPU進行加速并獲得了291倍的速度提升。

HAN：提供了靈活的接口可以將一個異構圖通過元路徑（metapath）轉變成同構圖。

Metapath2vec：新的元路徑采樣實現比原作者的實現快兩倍。

而且， DGL也發布了針對分子化學的模型庫DGL-Chem，以及訓練知識圖譜嵌入（Knowledge Graph Embedding）專用包DGL-KE。DGL-Chem提供了包括分子性質預測和分子結構生成等預訓練模型。

在單GPU上，DGL-KE能在7分鐘內使用經典的TransE模型訓練出FB15K的圖嵌入。相比GraphVite（v0.1.0）在4塊GPU上需要14分鐘。

DGL-KE首個版本發布了TransE、ComplEx和Distmult模型，支持CPU訓練、GPU訓練、CPU和GPU混合訓練以及單機多進程訓練。

DGL團隊表示，在未來也會發布更多模型和訓練方式。

傳送門

DGL異構圖教程：

https://docs.dgl.ai/tutorials/hetero/1_basics.html

DGL v0.4 更新筆記：

https://github.com/dmlc/dgl/releases

DGL-KE代碼及使用說明：

https://github.com/dmlc/dgl/tree/master/apps/kg

DGL-Chem 模型庫:

https://docs.dgl.ai/api/python/model_zoo.html#chemistry