北大、微軟提出NGra：高效大規(guī)模圖神經(jīng)網(wǎng)絡計算

目前，深度學習技術通常以深度神經(jīng)網(wǎng)絡（DNN）的形式展現(xiàn)。由于其在語音、視覺和自然語言處理等領域所取得的成功，深度學習技術越來越受歡迎。在這些領域中，底層數(shù)據(jù)表征的坐標通常具有規(guī)則的網(wǎng)格結構，這有利于包含大量類似于單指令多數(shù)據(jù)流（SIMD）數(shù)據(jù)并行運算的硬件加速機制（例如 GPU）。

目前，將深度學習模型應用在具有不規(guī)則圖結構的數(shù)據(jù)上成為了一種新興的趨勢 ，這種趨勢是由諸如社交網(wǎng)絡、知識圖譜、以及生物信息學和神經(jīng)科學（例如，蛋白質之間的交互或大腦中的神經(jīng)元連接）中的圖形的重要性所驅動的。這種趨勢也促使人們在其目標應用（例如，分類、嵌入、問答系統(tǒng)）上取得了當前最佳的結果。這些基于圖的神經(jīng)網(wǎng)絡（GNN）通常將神經(jīng)網(wǎng)絡模型應用在圖中與頂點和邊相關的特征上，傳播運算結果并進行聚合，從而生成下一級的特征。

現(xiàn)有的解決方案都不能很好地支持 GNN?，F(xiàn)有的圖處理引擎往往會提供一個類似于信息收集（Gather）——應用（Apply））——結果分發(fā)（Scatter）的 GAS 頂點程序模型，但這種方式無法在圖結構中表達和支持神經(jīng)網(wǎng)絡架構。TensorFlow、PyTorch、MxNet[8]、CNTK等深度學習框架旨在將神經(jīng)網(wǎng)絡表示為數(shù)據(jù)流圖，但并不直接支持圖傳播模型。此外，它們都不能提供處理大型圖所需的可擴展性，也不支持基于 GPU 的圖傳播 operator（將圖傳播轉換為稀疏操作）的高效實現(xiàn)。當前缺乏對這些需求的支持嚴重限制了充分挖掘大規(guī)模 GNN 潛力的能力，同時也為 DNN 與大型圖結構的結合在系統(tǒng)層面上提出了巨大的挑戰(zhàn)。

在本文中，作者介紹了首次實現(xiàn)支持大規(guī)模 GNN 的 NGra 系統(tǒng)，它從一個易于表達的編程模型發(fā)展到一個可擴展的、高效 GPU 并行處理引擎。NGra 將數(shù)據(jù)流與頂點程序的抽象自然地結合在了一個名為 SAGA-NN（Scatter-Apply Edge Gather-Apply Vertex with Neural Networks）的新模型中。SAGA 可以被認為是一個 GAS 模型的變體，SAGA-NN 模型中由用戶定義的函數(shù)使用戶能夠通過使用數(shù)據(jù)流抽象（而不是使用那些為處理傳統(tǒng)圖問題如 PageRank、連通分量、最短路徑而設計的算法）在頂點或邊數(shù)據(jù)（以張量形式被處理）上表達神經(jīng)網(wǎng)絡計算。

與在深度神經(jīng)網(wǎng)絡上的情況一樣，高效使用 GPU 對 GNN 的性能至關重要，而且由于處理大型圖結構的額外挑戰(zhàn)，對 GPU 的使用就顯得更為關鍵了。為了實現(xiàn)超越 GPU 物理限制的可擴展性，NGra 顯式地將圖結構劃分（頂點和邊數(shù)據(jù)）成塊，將一個在 SAGA-NN 模型中表達的 GNN 算法轉成了具有塊粒度 operator 的數(shù)據(jù)流圖，通過該數(shù)據(jù)流圖，作者可以在單個或多個 GPU 上進行基于塊的并行流處理。

然后，NGra 引擎的效率在很大程度上取決于 NGra 如何管理和調度并行流的處理過程，以及如何在 GPU 上實現(xiàn)關鍵的圖傳播 operator（結果分發(fā)和信息收集）。NGra 十分關注數(shù)據(jù)的位置，從而盡量減少 GPU 內(nèi)存內(nèi)外的數(shù)據(jù)交換，最大限度地提高 GPU 內(nèi)存中數(shù)據(jù)塊的重用性，同時以流形式同時進行數(shù)據(jù)遷移和計算。對于多 GPU 的情況，它使用了一種基于環(huán)的流機制，通過在 GPU 之間直接交換數(shù)據(jù)塊來避免主機內(nèi)存中的冗余數(shù)據(jù)遷移。SAGA-NN 模型中的結果分發(fā)和信息收集階段執(zhí)行沿著邊的頂點數(shù)據(jù)傳播，在稀疏結構上表現(xiàn)為矩陣乘法。眾所周知，在像 GPU 這樣數(shù)據(jù)并行的硬件上執(zhí)行稀疏的矩陣操作是非常困難的。因此，NGra 將圖傳播引擎支持的特殊 operator 引入到數(shù)據(jù)流圖中，并優(yōu)化其在 GPU 上的執(zhí)行。請注意，與其 基于 GPU 的圖引擎所關注的傳統(tǒng)圖處理場景不同，在 GNN 的場景下，由于每個頂點的數(shù)據(jù)可能就是一個特征向量，而不是簡單的標量，可變頂點數(shù)據(jù)本身可能無法被容納到 GPU 設備內(nèi)存中。因此，本文的方案更傾向于利用每個頂點數(shù)據(jù)訪問中的并行性來提高內(nèi)存訪問效率。

作者使用頂點程序抽象和用于圖傳播過程的自定義 operater 對 TensorFlow 進行擴展，從而實現(xiàn) NGra。結果證明 NGra 可以被擴展，然后通過利用單個服務器的主機內(nèi)存和 GPU 的計算能力，實現(xiàn)在（包含數(shù)百萬頂點和數(shù)百特征維度以及數(shù)億邊的）大型圖上對 GNN 算法的支持，而這不能直接通過使用現(xiàn)有的深度學習框架實現(xiàn)。與可以通過 GPU 支持的小型圖上的 TensorFlow 相比，NGra 可以取得大約 4 倍的運算加速。作者還對 NGra 中的多個優(yōu)化所帶來的性能提升進行了廣泛的評估，以證明它們的有效性。

圖 1：雙層 GNN 的前饋運算過程

圖 3：GNN 中每一層的 SAGA-NN 運算流程

NGra 系統(tǒng)

NGra 以用戶接口的形式提供了數(shù)據(jù)流和定點程序抽象的組合。在這種抽象之下，NGra 主要由以下幾部分組成：（1）一個將 SAGA-NN 模型中實現(xiàn)的算法轉換為塊粒度數(shù)據(jù)流圖的前端，它使得大型圖上的 GNN 計算可以在 GPU 中被實現(xiàn)；（2）一個制定最小化主機與 GPU 設備內(nèi)存之間數(shù)據(jù)遷移調度策略的優(yōu)化層，它能夠找到進行融合操作和去除冗余計算的機會；（3）一組高效的傳播操作內(nèi)核，它支持針對 GPU 中重復的數(shù)據(jù)遷移和計算的基于流的處理；（4）在運行時執(zhí)行數(shù)據(jù)流。由于 NGra 在很大程度上利用了現(xiàn)有的基于數(shù)據(jù)流的深度學習框架來實現(xiàn)在運行時執(zhí)行數(shù)據(jù)流，因此作者將重點放在本節(jié)前三個框架的設計上，因為它們是 NGra 系統(tǒng)的主要貢獻。

圖 4：用于 G-GCN 層上目標區(qū)間 V0 的基于塊的數(shù)據(jù)流圖。為了得到更清晰的可視化結果，在連接到 D2H 時，SAG 階段輸出張量置換被隱藏在了 SAG 的子圖中。

圖 13：TensorFlow（TF），cuSPARSE，以及 NGra（NG）在不同密度的圖上傳播內(nèi)核的時間。

圖 16：在大型圖上的不同應用中使用 NGra 的加速情況

論文：Towards Efficient Large-Scale Graph Neural Network Computing

論文地址：https://arxiv.org/abs/1810.08403

最近的深度學習模型已經(jīng)從低維的常規(guī)網(wǎng)格（如圖像、視頻和語音）發(fā)展到了高維的圖結構數(shù)據(jù)（如社交網(wǎng)絡、大腦連接和知識圖譜）上。這種演進催生了基于大型圖的不規(guī)則和稀疏模型，這些模型超出了現(xiàn)有深度學習框架的設計范圍。此外，這些模型不容易適應在并行硬件（如 GPU）上的有效大規(guī)模加速。

本文介紹了基于圖的深度神經(jīng)網(wǎng)絡（GNN）的第一個并行處理框架 NGra。NGra 提出了一種新的 SAGA-NN 模型，用于將深度神經(jīng)網(wǎng)絡表示為頂點程序，其每一層都具有定義好的（結果分發(fā)、應用邊、信息收集、應用頂點）圖操作階段。該模型不僅可以直觀地表示 GNN，而且便于映射到高效的數(shù)據(jù)流表示。NGra 通過自動圖形劃分和基于塊的流處理（使用 GPU 核心或多個 GPU），顯式地解決了可擴展性的挑戰(zhàn)，它仔細考慮了數(shù)據(jù)位置、并行處理和數(shù)據(jù)遷移的重疊等問題。NGra 雖然是稀疏的，但通過對 GPU 上的結果分發(fā)/信息收集 operator 進行高度優(yōu)化，進一步提高了效率。本文的評估結果表明，NGra 可以擴展到任何現(xiàn)有框架都無法直接處理的大型真實圖結構上，同時即使是在小規(guī)模的 TensorFlow 的多基線設計上，也可以實現(xiàn)高達 4 倍的加速。