亞馬遜＋紐約大學開源圖神經網路框架DGL：新手友好，與主流框架無縫銜接

科技 12-08

量子位授權轉載 | 公眾號 QbitAI

最近，紐約大學、紐約大學上海分校、AWS上海研究院以及AWS MXNet Science Team共同開源了一個面向圖神經網路及圖機器學習的全新框架，命名為Deep Graph Library（DGL）。

據介紹，這個框架在測試時表現出了良好的性能，能夠既快又好地部署深度神經網路。

開源具體信息已經在今天召開的NeurIPS 2018 ML Sys Workshop（Workshop on Systems for ML and Open Source Software）上公布，量子位編輯如下：

深度神經網路在近年來的發展和成功有目共睹。這些網路架構成功的一個重要原因在於其能有效地捕捉到了數據內在的相關性，並在不同的抽象級別構建表示(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的性能進行了測試：