亞馬遜馬超:如何使用DGL進行大規模圖神經網絡訓練?

演講嘉賓 | 馬超(亞馬遜應用科學家)

整理 | 劉靜

出品 | AI科技大本營(ID:rgznai100)

與傳統基於張量(Tensor)的神經網絡相比,圖神經網絡將圖 (Graph) 作為輸入,從圖結構中學習潛在的知識,該方法在近些年已被證明在許多場景可以取得很好的效果。然而,使用傳統的深度學習框架(比如 TensorFlow、Pytorch、MXNet)並不能方便地進行圖神經網絡的開發和訓練,而 DGL 作為專門面向圖神經網絡的框架,可以很好地彌補這一缺陷。該框架在開源後於國內外引起了強烈的反響。

基於此,在CSDN主辦的2019 AI開發者大會(AI ProCon 2019)上,亞馬遜應用科學家馬超,同時也是 DGL 項目的合作作者,發表了《使用 DGL 進行大規模圖神經網絡訓練》的主題演講。他從 API 使用、系統優化以及系統可擴展性等多個維度深入分享了 DGL 的設計思路。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

以下為馬超演講全文,AI科技大本營整理(ID:rgznai100):

現場大部分朋友都來自工業界,我想應該對傳統的圖算法瞭解更多,比如 PageRank 算法應該是我們學生時代學到的第一個傳統圖算法。但是突然之間,圖神經網絡算法火了起來。下面這張圖是近些年 GNN 論文發表數量的情況,我們可以看到從 2016 年開始,GNN 論文發表數量呈指數上升趨勢。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

什麼是圖神經網絡?

圖神經網絡在近年來的發展有目共睹。我們發現,光是去年的 ICML 上有近 50 篇論文都在講圖神經網絡。而今年的 KDD 和 NeurIPS 上也有非常多有意思的圖神經網絡研究。所以我們要問的第一個問題是:為什麼要在圖上做深度學習?這個問題其實很好回答:因為現實生活中很大一部分數據都可能是圖數據,而不是傳統深度學習算法裡的歐幾里得(Euclid)結構數據。

相比之下,我們做語音、圖像這類任務,輸入數據都是規整的張量,例如圖像中的像素一個個排列下來。但是現實生活中很多數據並不是這樣子,比如社交網絡。在社網絡中,如果每個用戶都被當作一個節點,用戶之間的關係被當作節點之間的邊,那麼整個數據就是一個圖,這是沒辦法用類似於像素這樣的歐幾里得結構來表示的;還有知識圖譜 (Knowledge Graph)這個例子。

比如,我們使用百度百科時,關鍵詞之間所組成關聯圖譜,也是圖結構;此外,還有一些研究人員在正在嘗試在醫療領域進行新藥或者以及新化學結構的預測和合成,而化學結構本質上也是圖。對整個圖而言,沒有一個節點的嵌入表示是孤立的,而傳統 CNN 和 RNN 分別把空間和時間上的相臨節點作為處理對象的方法並不能很好的刻畫出整個圖上的關係,侷限性很大。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

圖神經網絡能做什麼?

在聊 GNN 之前,我們首先思考一下大家熟悉的卷積神經網絡 CNN 在做什麼。比如指定一張圖片,CNN 幫助我們判斷是貓還是狗,本質上是卷積神經網絡學習了在某些圖像上的嵌入表示。而給定一張圖,如果用 GNN 來預測節點之間的關係,本質上是 GNN 學習了在圖上的嵌入表示。比如我們有一個 NBA 粉絲社群的關係圖,這時給你一個新的用戶,那麼即使我們本來不知道這個新用戶應該屬於哪個社群,但是仍然可以通過 GNN 來預測這個新用戶是不是勇士隊的球迷,因為 GNN 發現該用戶很多的關係節點都是勇士隊的球迷。

那麼 GNN 具體的工作機制是怎樣的?大家請看這張圖。在這張圖裡,不同顏色的圓圈代表圖上不同的節點。我們做圖神經網絡訓練時,每一個節點會給它相鄰的節點發消息 message,消息可以來自於邊上的特徵,也可以來自於節點上的特徵。鄰居節點拿到消息之後會進行聚合操作,然後通過一個神經網絡的非線性變換,例如 ReLu,並把它的輸出可以作為下一次發送消息的輸入,然後繼續發送給鄰居節點,通過這樣不斷迭代學習最終學會圖上的參數表示。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

相信這個過程會讓大家聯想起 PageRank 算法的流程,其實 PageRank 算法和這裡演示的 GNN 算法過程類似,只不過在 PageRank 裡沒有特徵向量和神經網絡的概念。

我們之前提到,從 2016 年開始出現了大量關於圖神經網絡的研究論文。但是即使你是一個圖神經網絡的專家,如果沒有好的工具,想要去復現已有的論文工作或者創造新的算法,仍然十分困難。比如遇到內存洩露或者內存爆掉的問題時,大部分研究人員都束手無策。除此之外,還有一個問題就是現在的深度學習框架,無論是Tensorflow、Pytorch 還是 MXNet,默認都是基於稠密的 Tensor 表示。

對於稠密張量的計算和自動求導,目前的框架已經可以做的很好了,而實現一個圖神經網絡程序卻不是遵循這樣的範式,大部分時間開發者都是在對稀疏數據進行操作。開發者從稠密數據的編程模式遷移到稀疏數據的編程模式上需要跨越巨大的 Gap。如果用戶現在想編寫圖神經網絡程序,可能需要花費很多時間進行程序上的優化,並且效果還不一定好。

正是因為存在這些問題,去年我們亞馬遜 AWS 聯合 NYU 以及 NYU上海,一起推出了 DGL 這個平臺,用來幫助用戶解決上述的這些問題。DGL 希望既能方便研究人員進行在圖神經網絡領域進行研究,又可以幫助工業界的用戶進行業務部署。

DGL:專門面向圖神經網絡的框架

如果大家有關注我們的官網的話(https://www.dgl.ai/),可以在上面找到 dgl 的源代碼、文檔教程,博客資源,論壇討論等等。我們的代碼目前已獲得 3000+ stars,450+ forks,60k+下載量,目前已經持續迭代到 0.4 版本。

DGL 發佈之後,用戶的主要構成來自中國和美國,這和我們團隊構成是一樣的。到目前為止,市面上也有了其他的圖神經網絡框架競爭者,那麼DGL的目標是什麼?其實主要是兩個方面:1. 我們要對已有深度學習框架的向前和向後兼容做到最好;2.在第一個目標基礎上儘可能提高框架的性能和可擴展性。

對於第一個目標,或者說是挑戰。在面對各種各樣的 GNN 算法需求是,DGL 框架能否提供一套靈活易用的 API 讓研究人員在框架上面可以快速不斷嘗試新的想法。經過大量研究現有的圖神經網絡論文以及代碼,DGL 最終選擇使用 message passing 作為底層的API 抽象。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

圖神經網絡框架最核心部分就是在處理 message passing,因此 DGL 把消息通信放在了最底層的API,用來適應廣泛的算法需求,對於大部分 GNN 算法而言,無論算法如何變化,通常都很難脫離 message passing 的表示。下面這張圖展示 DGL 為 message Passing API 提供了不同層級底層的封裝。

當然,光有 message passing 是不夠的,DGL 可以讓用戶實現自己需要的 UDF(用戶定義函數)。下圖是一個簡單的 GCN 的 demo,展示瞭如何用 DGL 編寫 GNN 程序。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

DGL 另一個重要的設計是,我們從一開始就把 DGL 設計成 “框架上的框架”。為了避免重複造輪子,DGL 在設計上採用了類似 Keras 的做法,基於已有的主流深度學習框架之上進行開發,目前 DGL 已經可以很好的支持 MXNet 和 Pytorch 兩個主流深度學習框架。

考慮到很大一部分用戶可能就是想利用現成的代碼把已有的算法拿來玩一玩,所以最近我們一口氣寫了很多現成的 NN moudule,方便大家可以直接調用。此外 DGL 還提供多種 GNN 模型的完整實現和運行腳本,歡迎大家都來玩一玩。包括我們在 0.4 版本即將推出的 Knowledge Graph model zoo.

DGL的性能優勢

DGL 在性能上做了很多系統優化。例如消息融合,DGL 通過稀疏矩陣乘法的轉化來做 message fusion。此外,為了讓系統在稀疏數據上的性能達到最好,DGL還專門啟動了一個名為 mini-Gun 的項目來做 kernel 優化。除此之外, DGL 內部還會通過 auto-batching 技術來加速批量圖的並行處理過程。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

談論了單機上的性能,那麼遇到大圖怎麼辦?比如微博有好幾億用戶,用戶的關係可能是百億級別,這樣的圖無論是對 GPU 還是單機來說都裝不下。因此 DGL 做了分佈式的架構。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?

這個架構有三層,最底層是採樣器,中間這層的 trainer 是訓練器,最上面是DGL-PS參數服務器。如果圖非常大,我們就先把大圖切分成很多小圖,再把它分到不同的機器或者不同的GPU上,每次訓練時只使用一小部分的 miini-batch 數據,並且通過參數服務器來存儲訓練過程中需要的模型參數,trainer 和參數服務器之間通過我們自己實現的網絡協議進行數據交互。

最後我想說的是,現在深度學習的主要數據結構是 Tensor,但是 Tensor 是真正的數據結構嗎?大家可以思考一下這個問題,如果把程序=算法+數據結構這句話擴展到深度學習領域中,會不會Graph就是深度學習的重要Date Structure?這也是為什麼說一定要抓住圖神經網絡這個機會。黑盒式的神經網絡一定不是長久走下去的道路,神經網絡的發展在某一天可能會遇到可解釋問題的瓶頸,圖神經網絡是否能為我提供一個神經網絡可解釋性的有效方法,值得我們思考。

亚马逊马超:如何使用DGL进行大规模图神经网络训练?
亚马逊马超:如何使用DGL进行大规模图神经网络训练?


分享到:


相關文章: