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

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

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

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

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

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

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

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

圖 1：雙層 GNN 的前饋運算過(guò)程

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

NGra 系統

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

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

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

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

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

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

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

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