NYU、AWS聯合推出：全新圖神經網路框架DGL正式發布

知識 12-09

機器之心發布

作者：DGL Team

深度神經網路在近年來的發展和成功有目共睹。這些網路架構成功的一個重要原因在於其能有效地捕捉到了數據內在的相關性，並在不同的抽象級別構建表示 (representation)。比如CNN刻畫了圖片中相鄰像素點間的空間不變性；RNN抓住了文本數據的有序線性結構。

CNN 和 RNN 的成功讓人思考——我們是否能將此思想拓展到其他結構的數據上呢？事實上，人們一直在探索如何將深度學習應用於更廣泛的結構數據中。早在 2014 年，Kai Sheng Tai 等人就研究了能在文本語法樹上訓練的樹神經網路模型 TreeLSTM。這個工作在一定程度上衝擊了大家用 RNN 處理文本的範式，並且用樹型結構看待文本數據開創了很多新的研究可能。從鏈表到樹，從樹到圖：近年來，對於圖神經網路（Graph Neural Network）的研究熱潮使得神經網路的應用領域大大增加。

小到分子結構圖（[Gilmer et al., 2017], [Jin et al., 2018]），大到知識圖譜（[Schlichtkrull et al., 2017]）、社交網路和推薦系統（[Ying et al., 2018]），圖神經網路都展現了強大的應用潛力。

廣義上說，沒有一個點的特徵表達是孤立的；豐富特徵的一個重要手段必然是融合該點和其他鄰居點的互動。從這個意義上來說，CNN 和 RNN 分別把空間和時間上的相鄰點看作互動對象是非常粗糙的。把圖和深度學習結合的意義正在於突破這個局限。

由於結構數據的普適性，Google, Deepmind, Pinterest 等公司也紛紛在這一領域發力。除了數據，深度神經網路模型的發展也有著更加動態和更加稀疏的趨勢。這源於目前深度神經網路對於算力的需求已經達到了一個驚人的地步。同時，我們觀察到現在還沒有一個框架既高效，又好用，能幫助開發新的模型。

（比如目前火熱的 BERT 模型在一台 8 塊 V100 的伺服器上預計訓練時間是 42 天。）

於是一個很現實的問題是如何設計「既快又好」的深度神經網路？也許更加動態和稀疏的模型會是答案所在。可見，不論是數據還是模型，「圖」應該成為一個核心概念。

基於這些思考，我們開發了 Deep Graph Library（DGL），一款面向圖神經網路以及圖機器學習的全新框架。

在設計上，DGL 秉承三項原則：

DGL 必須和目前的主流的深度學習框架（PyTorch、MXNet、TensorFlow等）無縫銜接。從而實現從傳統的 tensor 運算到圖運算的自由轉換。

DGL 應該提供最少的 API 以降低用戶的學習門檻。

在保證以上兩點的基礎之上，DGL 能高效並透明地並行圖上的計算，以及能很方便地擴展到巨圖上。

這裡，我們簡單介紹 DGL 是如何實現以上目標的。

設計一：DGL 是一個「框架上的框架」

神經網路計算的重要核心是稠密張量計算和自動求導。當前神經網路框架對此都已經有了非常良好的支持。此外，由於神經網路模型的日益模塊化，用戶對於這些深度學習的框架的黏性度也越來越高。

為了避免重複造輪子（build a wheel again），DGL 在設計上採取了類似 Keras 的做法——基於主流框架之上進行開發。但又不同於 Keras，DGL 並不限制用戶必須使用自己的語法。相反，DGL 鼓勵用戶在兩者間靈活地使用。比如用戶可以使用他們偏愛的框架編寫常見的卷積層和注意力層，而當遇到圖相關的計算時則可以切換為 DGL。用戶通過 DGL 調用的計算，經過系統優化，仍然調用底層框架的運算和自動求導等功能，因此開發和調試都能更快上手。用戶和 DGL 的交互主要通過自定義函數 UDF（user-defined function）。在下一節編程模型我們將繼續展開。目前 DGL 支持 Pytorch 以及 MXNet/Gluon 作為系統後端。

設計二：基於「消息傳遞」（message passing）編程模型

消息傳遞是圖計算的經典編程模型。原因在於圖上的計算往往可以表示成兩步：

發送節點根據自身的特徵（feature）計算需要向外分發的消息。

接受節點對收到的消息進行累和並更新自己的特徵。

比如常見的卷積圖神經網路 GCN（Graph Convolutional Network）可以用下圖中的消息傳遞模式進行表示（圖中 h 表示節點各自的特徵）。用戶可以定製化消息函數（message function），以及節點的累和更新函數（reduce function）來構造新的模型。事實上，在 Google 的 2017 年的論文中 [Gilmer et al. 2017] 將很多圖神經網路都歸納到了這一體系內。

DGL 的編程模型正基於此。以下是圖卷積網路在 DGL 中的實現（使用 Pytorch 後端）：

可以看到，用戶可以在整個程序中靈活地使用 Pytorch 運算。而 DGL 則只是輔助地提供了諸如 mailbox、send、recv 等形象的消息傳遞 API 來幫助用戶完成圖上的計算。在我們內部的項目預演中，我們發現熟悉深度學習框架的用戶可以很快上手 DGL 並著手開發新模型。

設計三：DGL 的自動批處理（auto-batching）

好的系統設計應該是簡單透明的。在提供用戶最大的自由度的同時，將系統優化最大程度地隱藏起來。圖計算的主要難點在於並行。

我們根據模型特點將問題劃分為三類。首先是處理單一靜態圖的模型（比如 citation graph），其重點是如何並行計算多個節點或者多條邊。DGL 通過分析圖結構能夠高效地將可以並行的節點分組，然後調用用戶自定義函數進行批處理。相比於現有的解決方案（比如 Dynet 和 TensorflowFold），DGL 能大大降低自動批處理的開銷從而大幅提高性能。

第二類是處理許多圖的模型（比如 module graph），其重點是如何並行不同圖間的計算。DGL 的解決方案是將多張圖合併為一張大圖的多個連通分量，從而將該類模型轉化為了第一類。

第三類是巨圖模型（比如 knowledge graph），對此，DGL 提供了高效的圖採樣介面，將巨圖變為小圖樣本，從而轉化為第一類問題。

目前 DGL 提供了 10 個示例模型，涵蓋了以上三種類別。其中除了 TreeLSTM，其餘都是 2017 年以後新鮮出爐的圖神經網路，其中包括幾個邏輯上相當複雜的生成模型（DGMG、JTNN）我們也嘗試用圖計算的方式重寫傳統模型比如 Capsue 和 Universal Transformer，讓模型簡單易懂，幫助進一步擴展思路。

我們也對 DGL 的性能進行了測試：