作者 | 付浩 憲鵬 祥洲等
美團搜索與NLP團隊在圖神經網絡的長期落地實踐中，基于業務實際場景，自主設計研發了圖神經網絡框架Tulong，以及配套的圖學習平臺，提升了模型的規模和迭代效率。
1. 前言
萬物之間皆有聯系。圖作為一種通用的數據結構，可以很好地描述實體與實體之間的關系。例如，在社交網絡中，用圖來表示用戶與用戶之間的好友關系；在電商網站中，用圖表示用戶與商品之間的點擊購買行為；在知識圖譜構建中，還可以用圖表示實體與實體間多樣的關系。另一方面，深度學習技術在計算機視覺、自然語言處理、語音處理等領域均已取得了巨大的成功。深度學習技術將圖像、文本、語音等多種多樣的數據轉化為稠密的向量表示，提供了表示數據的另一種方式。借助于硬件日益強大的計算能力，深度學習可以從海量數據中學習到數據之間復雜多樣的相關性。
這會讓人不禁思考，深度學習能否應用到更廣闊的領域，比如——圖？事實上，早在深度學習興起之前，業界就已經開始了圖嵌入(Graph Embedding)技術的探索[1]。早期的圖嵌入算法多以啟發式的矩陣分解、概率圖模型為主；隨后出現了以DeepWalk[2]和Node2vec[3]為代表的、較為“淺層”的神經網絡模型；最后，以GCN[4]為代表的一系列研究工作，打通了圖信號處理與神經網絡之間的壁壘，奠定了當前基于消息傳遞機制的圖神經網絡(GNN: Graph Neural Network)模型的基本范式。
近年來，圖神經網絡逐漸成為學術界的研究熱點之一[5]。在工業界，圖神經網絡在電商搜索、推薦、在線廣告、金融風控、交通預估等領域也有諸多的落地應用，并帶來了顯著收益。
由于圖數據特有的稀疏性（圖的所有節點對之間只有少量邊相連），直接使用通用的深度學習框架（例如TensorFlow和PyTorch）訓練往往性能不佳。工欲善其事，必先利其器。針對圖神經網絡的深度學習框架應運而出：PyG (PyTorch Geometric)[6]和DGL (Deep Graph Library)[7]等開源框架大幅提升了圖神經網絡的訓練速度，并且降低了資源消耗[17][18]，擁有活躍的社區支持。很多公司根據自身業務特點，也紛紛建設自有的圖神經網絡框架。美團搜索與NLP團隊在長期的落地實踐中，總結實踐經驗，在訓練的規模和性能、功能的豐富性、易用性等方面進行了大量優化。本文首先介紹我們在過往落地應用中遇到的實際問題和挑戰，然后再介紹具體的解決方案。
1.1 問題和挑戰
從工業界落地應用的角度來看，一個“好用”的圖神經網絡框架至少具備以下特點。
（1）完善支持當前流行的圖神經網絡模型。
從圖本身的類型來看，圖神經網絡模型可以分為同質圖(Homogeneous Graph)、異質圖(Heterogeneous Graph)、動態圖(Dynamic Graph)等類型。從訓練方式來看，又可以分為全圖消息傳遞[4]和基于子圖采樣的消息傳遞[8]等類型。從推理方式來看，還可以分為直推式和歸納式[9]。
除此之外，下游任務除了經典的節點分類、鏈接預測和圖分類，還有許多領域相關端到端的預測任務。在實際應用中，不同業務場景對圖神經網絡的模型和下游任務的需求是不同的，需要個性化定制。例如在美食推薦場景中，存在用戶、商家、菜品等節點，刻畫其相互關系可以用同質圖或異質圖；為了刻畫用戶在不同時間的偏好，可能還需要使用動態圖模型；針對推薦系統的召回和排序兩個階段，還需要設計不同的訓練任務。盡管現有框架都提供常見模型的實現，但簡單調用這些模型不能滿足上述需求。此時便需要用戶自行開發模型和訓練流程代碼，這就帶來了額外的工作量。如何幫助用戶更便捷地實現定制模型是一個不小的挑戰。
（2）以合理的代價支持大規模圖上的模型訓練。
在業務落地應用中，圖的規模往往很大，可以達到數十億甚至數百億條邊。我們在初期的嘗試中發現，使用現有框架，只能在分布式環境下訓練百億邊規模的模型，消耗較多的硬件資源（數千CPU和數TB內存）。我們希望單機即可在合理的時間內訓練百億邊規模的模型，從而降低對硬件資源的需求。
（3）與業務系統無縫對接。
圖神經網絡的完整落地流程至少包括：基于業務數據構圖、離線訓練和評測模型、線上推理、業務指標觀測等步驟。要讓圖神經網絡技術成功落地應用，需要充分理解業務邏輯和業務需求，統一并高效地管理業務場景。同樣以美食推薦場景為例，線上日志記錄了曝光、點擊、下單等行為事件，知識圖譜提供了商家和菜品豐富的屬性數據，如何從這些異質的數據構造圖，要結合業務實際多次實驗確定。合適的工具能提升對接業務數據的效率，然而現有的圖神經網絡框架大多聚焦在模型的離線訓練和評測，缺乏此類工具。
（4）研發人員易于上手，同時提供充足的可擴展性。
從研發效率的角度來說，自建圖神經網絡框架的目的是減少建模中的重復工作，讓研發人員的精力集中在業務本身的特性上。因此，一個“好用”的圖神經網絡框架應當易于上手，通過簡單地配置即能完成多數任務。在此基礎上，對于一些特殊的建模需求，也能提供適當的支持。
1.2 美團的解決方案
美團搜索與NLP團隊在搜索、推薦、廣告、配送等業務的長期落地實踐中，總結實踐經驗，自主設計研發了圖神經網絡框架Tulong以及配套的圖學習平臺，較好地解決了上述問題。
- 首先，我們對當前流行的圖神經網絡模型進行了細粒度的剖析，歸納總結出了一系列子操作，實現了一套通用的模型框架。簡單修改配置即可實現許多現有的圖神經網絡模型。
- 針對基于子圖采樣的訓練方式，我們開發了圖計算庫“MTGraph”，大幅優化了圖數據的內存占用和子圖采樣速度。單機環境下，相較于DGL訓練速度提升約4倍，內存占用降低約60%。單機即可實現十億節點百億邊規模的訓練。
- 圍繞圖神經網絡框架Tulong，我們構建了一站式的圖學習平臺，為研發人員提供包括業務數據接入、圖數據構建和管理、模型的訓練和評測、模型導出上線等全流程的圖形化工具。
- Tulong實現了高度可配置化的訓練和評測，從參數初始化到學習率，從模型結構到損失函數類型，都可以通過一套配置文件來控制。針對業務應用的常見場景，我們總結了若干訓練模版，研發人員通過修改配置即可適配多數業務場景。例如，許多業務存在午晚高峰的周期性波動，我們為此設計了周期性動態圖的訓練模板，可以為一天中不同時段產生不同的GNN表示。在美團配送業務的應用中，需要為每個區域產生不同時段下的GNN表示，作為下游預測任務的輸入特征。開發過程中，從開始修改配置到產出初版模型僅花費三天；而在此之前，自行實現類似模型方案花費約兩周時間。
2. 系統概覽
如下圖1所示，Tulong配套圖計算庫和圖學習平臺構成了一套完整系統。系統自底向上可以分為以下3個組件。
圖1 圖神經網絡計算引擎、框架和平臺的系統架構
（1）圖以及深度學習引擎
我們把圖神經網絡的底層算子分為三類：圖結構查詢、稀疏張量計算和稠密張量計算。我們開發了圖計算庫MTGraph提供圖數據的存儲和查詢功能，深度優化了內存占用和子圖采樣速度。MTGraph兼容PyTorch和DGL，用戶可以在MTGraph的基礎上直接編寫基于DGL的模型代碼。
（2）Tulong框架
Tulong框架首先封裝實現了訓練圖神經網絡所需的基本組件，包括圖和特征數據的預處理流程、子圖采樣器、通用的GNN模型框架，以及包括訓練和評測在內的基礎任務。基于上述組件，Tulong框架提供豐富的預定義模型和訓練/推理流程，用戶通過修改配置文件即可在業務數據上訓練和評測GNN模型。
（3）圖學習平臺
圖學習平臺旨在簡化離線的模型開發和迭代過程，同時簡化業務系統的對接流程。圖學習平臺提供一系列的可視化工具，簡化從業務數據接入到模型上線的全流程。下文將從模型框架、訓練流程框架、性能優化和圖學習平臺等四個方面詳細介紹各個模塊的分析和設計方案。
3. 模型框架
我們從工程實現的角度，歸納總結了當前主流圖神經網絡模型的基本范式，實現一套通用框架，以期涵蓋多種GNN模型。以下按照圖的類型（同質圖、異質圖和動態圖）分別討論。
3.1 同質圖
同質圖(Homogeneous Graph)可以定義為節點集合和邊集合：，一條邊表示節點與節點相連。節點和邊上往往還附加有特征，我們記為節點的特征，為邊的特征。包括PyG和DGL在內的許多圖神經網絡框架，都對同質圖上的GNN進行過歸納，提出了相應的計算范式。例如，DGL把GNN的前向計算過程歸納為消息函數(message function)、聚合函數(reduce function)和更新函數(update function)[7]。
我們擴展了聚合函數的種類，提出一種更加通用的計算范式：
上述計算范式仍然分為生成消息、聚合消息、更新當前節點三個步驟，具體包括：
- 層次維度的聚合函數：用于聚合同一節點在模型不同層次的表示。例如，多數GNN模型中，層次維度的聚合函數為上一層的節點表示；而在JKNet[10]中，層次維度的聚合函數可以設定為LSTM[11]。
- 消息函數：結合起始節點和目標節點，以及邊的特征，生成用于消息傳遞的消息向量。
- 節點維度的聚合函數：匯集了來自鄰居節點的所有消息向量。值得注意的是，也可以有不同的實現。例如，在GCN中為所有鄰居節點，而在GraphSage[9]中為鄰居節點的子集。
- 更新函數：用于聚合節點自身在上一層和當前層的表示。
不難看出，上述計算范式可以覆蓋當前大多數GNN模型。在工程實踐中，我們將上述函數進一步分拆細化，預先提供了多種高效的實現。通過配置選項即可實現不同的組合搭配，從而實現多數主流的GNN模型。
3.2 異質圖
相比于同質圖，異質圖(Heterogeneous Graph)擴充了節點類型和邊類型。比如，學術引用網絡[13]中包含論文、作者、機構等類型的節點，節點直接通過“論文引用其他論文”、“作者撰寫論文”、“作者屬于機構”等類型的邊相連，如下圖2所示：
圖2 同質圖與異質圖的比較
我們把異質圖視為多個二分圖的疊加，每一個二分圖對應于一種邊類型。上述的學術引用網絡可以表示成“論文-引用-論文”、“作者-撰寫-論文”、“作者-屬于-機構”，共計三個二分圖，同質圖的GNN模型框架稍加修改即可在二分圖上應用。
在此基礎上，一個節點在不同的二分圖中會產生不同的表示。我們進一步提出邊類型維度的聚合函數，用于聚合節點在不同二分圖中的表示（如下圖3所示）。框架中同樣提供邊類型緯度聚合函數的多種實現，可以通過配置選項調用。例如，要實現RGCN，可以在二分圖上應用GCN，然后在邊類型維度上取平均。
圖3 異質圖模型框架
3.3 動態圖
動態圖(Dynamic Graph)是指隨時間變化的圖。與之相對的，上述的同質圖和異質圖可以稱為靜態圖。比如，學術引用網絡會隨時間不斷擴張，用戶與商品的交互圖會隨用戶興趣而變化。動態圖上的GNN模型旨在生成給定時間下的節點表示。根據時間粒度的粗細，動態圖可分為離散時間動態圖和連續時間動態圖。
在離散時間動態圖中，時間被劃分為多個時間片（例如以天/小時劃分），每個時間片對應一個靜態的圖。離散時間動態圖的GNN模型通常在每個時間片上單獨應用GNN模型，然后聚合節點在不同時間的表征[14]。我們把聚合過程抽象為離散時間維度的聚合函數，同樣提供預定義的實現。此外，Tulong框架還提供離散時間動態圖數據的加載和管理機制，僅在內存中保留必須的時間片，降低硬件資源的消耗。
圖4 離散時間動態圖GNN模型框架在連續時間動態圖中，每條邊附有時間戳，表示交互事件發生的時刻。相比于靜態圖，連續時間動態圖中的消息函數還依賴于給定樣本的時間戳以及邊的時間戳。此外，鄰居節點必須與時間有關，例如鄰居節點中不能出現時刻之后才出現的節點。針對此問題，我們開發了多種連續時間動態圖上的鄰居節點采樣器，可以在指定的時間范圍內，高效地采樣鄰居節點。
圖5 連續時間動態圖GNN模型框架
以上分析了同質圖、異質圖和動態圖的計算范式，我們從中抽取出通用的函數（算子），包括消息函數、聚合函數、更新函數、鄰居節點函數，并給出多種預定義的實現。框架用戶通過配置選項即可拼裝組合算子，從而實現需要的GNN模型。
4. 訓練流程框架
訓練GNN模型通常包括加載數據、定義GNN模型、訓練和評測、導出模型等流程。由于GNN模型和訓練任務的多樣性，在實際開發過程中，用戶往往要針對自己的場景自行編寫模型和流程代碼，處理繁瑣的底層細節讓用戶難以集中到算法模型本身的調優上。GraphGym[12]和DGL-Go[16]試圖解決這一問題，通過集成多種模型和訓練任務，同時簡化接口，可以讓用戶較為直接地上手和訓練GNN模型。
我們通過更加“工業化”的方式解決這一問題（如下圖6所示），框架被分為兩層：基礎組件和流程組件。基礎組件聚焦于單一的功能，例如圖數據組件只維護內存中的圖數據結構，不提供圖上的采樣或張量計算功能；圖上的采樣功能通過圖采樣器來提供。流程組件通過組裝基礎組件提供較為完整的數據預處理、訓練和評測流程，例如訓練流程組合了圖數據、圖采樣器、GNN模型等組件，提供完整的訓練功能。
圖6 訓練流程框架
更上一層，我們提供多種流程配置模板和GNN模型模板。模板對外暴露若干超參，例如訓練數據路徑、模型類型、學習率等參數，結合用戶指定的超參后就可以完整定義一次訓練任務。換言之，基于模板和參數即可完整復現一次GNN模型實驗。框架將會解析這些配置，并生成可執行的應用。
舉例來說，用戶可以選擇GraphSage模型的配置模板，以及鏈接預測任務的訓練模板，指定模型層數和維度，以及訓練評測數據路徑，即可開始訓練基于GraphSage的鏈接預測模型。
5. 性能優化
隨著業務的發展，業務場景下圖的規模也愈發龐大。如何以合理的代價，高效訓練數十億乃至百億邊規模的GNN模型成為亟需解決的問題。我們通過優化單機的內存占用，以及優化子圖采樣算法，來解決這一問題。
5.1 圖數據結構優化
圖數據結構的內存占用是制約可訓練圖規模的重要因素。以MAG240M-LSC數據集[13]為例，添加反向邊后圖中共有2.4億節點和35億邊。在基于子圖采樣的訓練方式下，PyG和DGL單機的圖數據結構均需要占用100GB以上的內存，其它開源框架的內存占用往往更多。在更大規模的業務場景圖上，內存占用往往會超出硬件配置。我們設計實現了更為緊湊的圖數據結構，提升了單機可承載的圖規模。我們借助圖壓縮技術降低內存占用。不同于常規的圖壓縮問題，GNN的場景下需要支持隨機查詢操作。例如，查詢給定節點的鄰居節點；判斷給定的兩個節點在圖中是否相連。我們對此提出的解決方案包括兩部分：
- 圖數據預處理和壓縮：首先分析圖的統計特征，以輕量級的方式對節點進行聚類和重新編號，以期讓編號接近的節點在領域結構上也更為相似。隨后調整邊的順序，對邊數據進行分塊和編碼，產生“節點-分塊索引-鄰接邊”層次的圖數據文件（如下圖7所示）。最后，如果數據包含節點特征或邊特征，還需要將特征與壓縮后的圖對齊。
圖7 壓縮后的圖數據結構
- 圖的隨機查詢：查詢操作分為兩步：首先定位所需的邊數據塊，然后在內存中解壓數據塊，讀取所查詢的數據。例如在查詢節點和是否相連時，首先根據兩個節點的編號計算邊數據塊的地址，解壓數據塊后獲得少量候選鄰接邊（通常不多于16條），然后查找是否包含邊。
經過壓縮，加載MAG240M-LSC數據集僅需15GB內存。百億乃至千億邊規模圖的內存占用顯著降低，達到單機可承載的程度，如下圖8所示：
圖8 圖數據結構內存占用對比
5.2 子圖采樣優化
子圖采樣是GNN模型訓練的性能瓶頸之一。我們發現在某些業務圖中，子圖采樣的耗時甚至占訓練整體的80%以上。我們分別針對靜態圖和動態圖，設計實現了多種高效的鄰居節點采樣算法。主要的優化手段包括：
- 隨機數發生器：相比于通信加密等應用，圖上的采樣對于隨機數發生器的“隨機性”并沒有苛刻的要求。我們適當放松了對隨機性的要求，設計實現了更快速的隨機數發生器，可以直接應用在有放回和無放回的采樣操作中。
- 概率量化：有權重的采樣中，在可接受的精度損失下，將浮點數表示的概率值量化為更為緊湊的整型。不僅降低了采樣器的內存消耗，也可以將部分浮點數操作轉化為整型操作。
- 時間戳索引：動態圖的子圖采樣操作要求限定邊的時間范圍。采樣器首先對邊上的時間戳構建索引，采樣時先根據索引確定可采樣邊的范圍，然后再執行實際的采樣操作。
經過以上優化，子圖采樣速度相較于DGL取得了2到4倍的提升（如下圖9所示）。某業務場景圖A（2億節點40億邊）使用DGL訓練耗時2.5小時/epoch，經過優化可達0.5小時/epoch。某業務場景圖B（2.5億節點124億邊）原本只能分布式訓練，耗時6小時/epoch；經過優化，單機即可訓練，速度可達2小時/epoch。
圖9 子圖采樣速度對比（2層，每層20條鄰接邊）
6. 圖學習平臺
圖學習平臺旨在簡化離線的模型開發迭代過程，同時簡化業務系統的對接流程。一個完整的模型開發迭代過程至少包括三個階段：準備數據集、定義模型和訓練任務、訓練和評測模型。我們分析用戶在這三個階段的需求，提供相應工具提升開發效率：
- 數據集管理：從業務數據構造圖是模型開發的第一步，圖學習平臺提供基于Spark的構圖功能，可以將Hive中存儲的業務數據轉化為Tulong自定義的圖數據格式。業務數據經常以事件日志的方式存儲，如何從中抽象出圖，有大量的選擇。例如，在推薦場景中，業務日志包含用戶對商家的點擊和下單記錄，除了把"用戶-點擊-商家"的事件刻畫為圖以外，還可以考慮刻畫短時間內共同點擊商家的關系。除此之外，還可以引入額外的數據，比如商家的地理位置、商家在售的菜品等。究竟使用何種構圖方案，需要經過實驗才能確定。對此，圖學習平臺提供了圖形化的構圖工具（如下圖10所示），幫助用戶梳理構圖方案；同時還提供圖數據集的版本管理，方便比較不同構圖方案的效果。
圖10 圖形化的構圖工具
- 實驗管理：確定圖數據之后，建模方案和訓練策略是影響最終效果的關鍵。例如，應該用何種GNN模型？損失函數如何選取？模型超參和訓練超參如何確定？這些問題也需要經過大量實驗才能回答。基于Tulong框架，建模方案和訓練策略可以通過一組配置來控制。圖學習平臺提供配置的可視化編輯器和版本管理功能，方便比較不同的方案的優劣。
- 流程管理：有了圖數據集和建模/訓練方案后，還需要讓整個流程自動化。這是模型上線的必要條件，同時也有利于團隊成員復現彼此的方案。圖學習平臺針對常見的“構圖、訓練、評測、導出”流程提供了自動化的調度，在適當的時候可以復用前一階段的結果，以提升效率。例如，如果數據集的定義沒有變化，可以跳過Spark構圖階段直接使用已有的圖數據。此外，針對模型上線的需求，平臺提供構圖和建模方案整合和定時調度等功能。?
7. 總結
本文介紹了美團搜索與NLP團隊在圖神經網絡框架建設方面的實踐經驗，包括GNN模型歸納抽象、基本框架、性能優化，以及上層工具等方面的思考和關鍵設計。框架的設計思路來源于業務落地所遇到的實際問題，例如針對大規模圖的優化、多人協作中的流程管理等；同時也吸收借鑒了學術界的最新研究進展，例如動態圖的計算范式等。除了技術層面的優化，框架的建設也得益于工程團隊和算法團隊的緊密配合，基于共同的、有深度的認知才得以讓項目順利推進。
借助于Tulong框架，圖神經網絡技術已在美團搜索、推薦、廣告、配送多個業務場景落地應用，并取得了較為可觀的業務收益。我們相信圖神經網絡還有更加廣闊的應用前景，作為基礎設施的圖神經網絡框架也值得繼續優化完善。
8. 作者簡介
付浩、憲鵬、祥洲、玉基、徐灝、夢迪、武威等，均來自美團平臺/搜索與NLP部。
