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