效能提升19倍，DGL重大更新支援億級規模圖神經網路訓練

神經網路 · 發表 2019-05-12 13:15:13

摘要：我們在去年12月釋出了Deep Graph Library (DGL)的首個公開版本。在過去的幾個版本的更新中，DGL主要注重框架的易用性，比如怎樣設計一系列靈活易用的介面，如何便於大家實現各式各樣的圖神經網路（GNN）模型，以及怎樣和主流深度學習框架（如PyTorch，MXNet等）整合。...

我們在去年12月釋出了Deep Graph Library (DGL)的首個公開版本。在過去的幾個版本的更新中，DGL主要注重框架的易用性，比如怎樣設計一系列靈活易用的介面，如何便於大家實現各式各樣的圖神經網路（GNN）模型，以及怎樣和主流深度學習框架（如PyTorch，MXNet等）整合。因為這些設計，讓DGL快速地獲得了社群的認可和接受。然而天下沒有免費的午餐，不同的框架對於相同的運算支援程度不同，並且普遍缺乏圖層面上的計算原語，導致了計算速度上的不足。隨著DGL介面的逐漸穩定，我們終於可以騰出手來解決效能問題。即將釋出的DGL v0.3版本中，效能問題將得到全面而系統地改善。

相比當前的DGL穩定版本v0.2，DGL v0.3在效能上取得了顯著提升。相比v0.2, DGL v0.3訓練速度提高了19倍，並且大幅度降低了記憶體使用量，使得單GPU上能訓練的圖的大小提高到原來的8倍。比起PyG等其他框架，DGL不但訓練更快，而且能夠在巨大的圖上（5億節點，250億邊）訓練圖神經網路。

接下來，我們將介紹DGL v0.3的重要特性之一 — 訊息融合（Fused Message Passing)。我們會逐一解釋，為什麼普通的訊息傳遞無法拓展到大圖上以及訊息融合是怎麼解決這一問題的。更多細節可以參考我們被 ICLR’19 的 RLGM workshop 所收錄的論文[1]。

大圖訓練的效能瓶頸

絕大多數圖神經網路模型遵循訊息傳遞的計算正規化，使用者需要提供兩個函式：

訊息函式：在邊上觸發，定義瞭如何計算髮送給相鄰節點的訊息。
累和函式：在點上觸發，定義瞭如果在點上累和收到的訊息。

下圖中，使用者自定義的訊息函式用表示。訊息函式將點 i 和 j 上的特徵，以及邊i->j上的特徵作為輸入，生成邊上的訊息（黃色方框）。在每個節點上，使用者定義的累和函式將訊息累和，然後呼叫另一個使用者定義的更新函式更新節點的特徵。

普通的訊息傳遞很容易在DGL中實現：首先，我們通過 send 介面呼叫訊息函式，然後通過recv 介面呼叫累和函式。下面的例子實現了目前流行的圖卷積網路Graph Convolution Network（GCN)。

# 使用自定義訊息函式和累和函式計算圖卷積
G.update_all(lambda edges: {'m' : edges.src['h']}, 
lambda nodes: {'h' : sum(nodes.mailbox['m'], axis=1)})

以上的程式碼非常簡潔易懂，但效能卻不佳。原因在於訊息傳遞的過程中實際生成了訊息張量（message tensor）。訊息張量的大小正比於圖中邊的數量，因而當圖增大時，訊息張量消耗的記憶體空間也會顯著上升。以 GraphSage 論文中的 Reddit 資料集（23.2萬節點，1.14億邊）為例，如果我們用上述程式碼訓練 GCN ，點上的特徵會被拷貝成邊上的資訊，這會導致記憶體使用量驟增500倍。除了浪費記憶體，該做法還使得訪存變得更為頻繁，進而導致 GPU 的利用率降低。

訊息融合解決大圖訓練難題

為了避免生成訊息張量帶來的額外開銷，DGL實現了訊息融合技術。DGL將 send 和 recv 介面合併成 send_and_recv（見下圖）。DGL的後端通過自己的CUDA程式碼，在每個GPU執行緒中將源節點特徵載入其本地記憶體並計算訊息函式，然後將計算結果直接累和到目標節點，從而避免生成訊息張量。

為實現訊息融合，DGL提供了一系列預先定義好的內建函式。儘管這限制了使用者對訊息函式和累和函式的選擇，但DGL提供了非常豐富的內建函式以實現絕大多數GNN模型。當然，使用者也可以選擇自己定義訊息函式和累和函式，這種情況下，DGL不會進行訊息融合優化。

另外在反向傳播中，由於訊息張量沒有儲存，因此需要被重新計算。實際操作中，許多訊息函式的求導都不需要使用到訊息張量（比如拷貝源節點特徵到邊上），而我們的實現也利用了這一特性。

在DGL中使用訊息融合

使用訊息融合非常簡單。比如，我們可以用copy_src內建訊息函式和sum內建累和函式改寫先前的GCN實現：

import dgl.function as fn
G = ... # 任意圖結構
# 將源節點的特徵h拷貝為訊息，並在目標節點累和生成新的特徵h。
G.update_all(fn.copy_src('h', 'm'), fn.sum('m', 'h'))

圖注意力模型 Graph Attention Network (GAT) 則可以用 src_mul_edge 內建訊息函式和 sum內建累和函式組合實現：

# 這裡假設注意力分數為邊上特徵e
G.update_all(fn.src_mul_edge('h', 'e', 'm'), fn.sum('m', 'h'))

DGL v0.3 將支援以下內建函式：

訊息函式可以是從源節點特徵、邊特徵、目標節點特徵三者中選任意兩個進行加、減、乘、除運算。
DGL支援特徵維度上的廣播語義（broadcasting semantics）。這在多頭注意力模組中非常常見。
累和函式可以是sum, max, min, prod。

我們推薦使用者儘可能多的使用DGL的內建函式來定義圖神經網路，這樣DGL可以利用訊息融合來提高效能。雖然這在上手上會有些門檻，但它對效能的提升是非常顯著的（詳見下一章節）。

效能測試

為了理解訊息傳遞融合帶來的效能提升，我們對DGL v0.3和DGL v0.2以及PyG（Pytorch Geometric v1.2.0）進行比較。其中PyG使用了普通的訊息傳遞實現，因此在整個過程中會生成訊息張量。

我們首先在主流的資料集上測試了GCN和GAT模型的效能，所有的實驗使用了模型論文中的引數設定。實驗在AWSp3.2xlarge instance上進行，該機器配備有NVIDIA V100 GPU （16GB 視訊記憶體）。

從表中可見，即將釋出的DGL v0.3在效能上有顯著提升，尤其在GAT模型上，訓練速度提升了19倍，而這都是因為使用了訊息融合技術。在小圖上（比如Cora，CiteSeer和PubMed），訓練的計算量和記憶體使用量幾乎不隨圖的大小發生變化，和PyG相比，DGL有微小且固定的額外開銷。然而，當在相對較大的圖上（比如從Reddit抽取出來的圖）訓練時，PyG很快便耗盡了記憶體，而DGL則可以輕鬆地將資料儲存在GPU上進行計算。

我們使用合成的圖進一步分析DGL的效能：

我們首先固定圖的密度（0.0008），通過調節圖的節點數來觀察GCN和GAT的訓練速度。從圖中可見，DGL可以在多達50萬節點的圖上訓練GCN模型，比PyG的最大容量高出一倍。此外，DGL的訓練速度比PyG快了3.4倍。

然後我們固定圖的節點數，通過調節圖的密度來觀察訓練速度。對GCN和GAT模型，相較PyG，DGL可以支援8倍多的邊，並且訓練快7.5倍。

我們還在一箇中等大小的圖上（3.2萬節點，密度0.0008）通過調節隱含層的大小來觀察訓練速度。對於GCN模型，儘管PyG能夠支1024個隱含單元，但其訓練速度比DGL慢了4倍。對於GAT模型，PyG最多隻能支援32個隱含單元，而DGL可以支援到256個。

最後，我們想測試DGL的效能極限，瞭解DGL在單機情況下能夠支援的最大的圖的規模。我們在AWSx1.32xlarge （2TB 記憶體）上用CPU訓練GCN。實驗表明，DGL可以支援到5億節點250億邊的圖。

接下來期待什麼

DGL團隊正在積極開發其設計路線圖上的功能特性。實際上，DGL專案開始之初，團隊成員就考慮到了絕大多數效能優化，比如，DGL一直提倡使用其內建函式而非自定義函式，儘管內建函式只有在訊息融合時才能發揮出其優勢。以下是DGL團隊正在積極拓展的方向：

撰寫更詳細的部落格介紹如何在算力強大的CPU機器上覆現大圖的實驗結果。
支援異構圖結構。
用GPU加速圖上的遍歷和訪問。

DGL一直努力接近使用者和社群，並且渴望得到使用者的反饋。如果您想要儘早嘗試即將在v0.3版本釋出的新特性，可以克隆DGL的GitHub倉庫，切換到kernel分支，然後從原始碼編譯DGL專案。

1.https://rlgm.github.io/papers/49.pdf