項(xiàng)目簡(jiǎn)介

本項(xiàng)目來源于飛槳AI for Science共創(chuàng)計(jì)劃的論文復(fù)現(xiàn)賽題，復(fù)現(xiàn)論文為《AMGNET: multi-scale graph neural networks for flow field prediction》。該論文主要采用圖神經(jīng)網(wǎng)絡(luò)，因?yàn)樵谟?jì)算流體力學(xué)中計(jì)算域被網(wǎng)格離散化，這與圖結(jié)構(gòu)天然契合。論文中通過訓(xùn)練 CFD 仿真數(shù)據(jù)，構(gòu)建一種數(shù)據(jù)驅(qū)動(dòng)模型進(jìn)行流場(chǎng)預(yù)測(cè)。本文將與大家分享基于飛槳完成該論文的復(fù)現(xiàn)過程，歡迎大家一起溝通學(xué)習(xí)。

環(huán)境依賴

01 硬件與框架

本項(xiàng)目需要的硬件環(huán)境與框架要求如下所示：

• GPU Memory >= 8GB

• 飛槳== 2.4.0

• Python ==3.7.4

• PGL == 2.2.4

• Matplotlib == 3.5.3

• pyamg == 4.2.3

• scipy

02 本地安裝

通過以下指令完成飛槳以及各個(gè)庫的安裝。

模型介紹

01 模型框架

AMGNET 是一種基于圖神經(jīng)網(wǎng)絡(luò)的 CFD 計(jì)算模型，該模型可以預(yù)測(cè)在不同物理參數(shù)下的流場(chǎng)。

圖3.1為模型的網(wǎng)絡(luò)結(jié)構(gòu)圖，該模型的基本原理就是將網(wǎng)格結(jié)構(gòu)轉(zhuǎn)化為圖結(jié)構(gòu)，然后通過網(wǎng)格中節(jié)點(diǎn)的物理信息、位置信息以及節(jié)點(diǎn)類型對(duì)圖中的節(jié)點(diǎn)和邊進(jìn)行編碼。接著對(duì)得到的圖神經(jīng)網(wǎng)絡(luò)使用基于代數(shù)多重網(wǎng)格算法(?Olson and Schroder, 2018?)的粗化層進(jìn)行粗化，將所有節(jié)點(diǎn)分類為粗節(jié)點(diǎn)集和細(xì)節(jié)點(diǎn)集，其中粗節(jié)點(diǎn)集是細(xì)節(jié)點(diǎn)集的子集。粗圖的節(jié)點(diǎn)集合就是粗節(jié)點(diǎn)集，于是完成了圖的粗化，縮小了圖的規(guī)模。粗化完成后通過設(shè)計(jì)的圖神經(jīng)網(wǎng)絡(luò)信息傳遞塊( GN )來總結(jié)和提取圖的特征。之后圖恢復(fù)層采用反向操作，使用空間插值法(?Qi et al.,2017?)對(duì)圖進(jìn)行上采樣。例如要對(duì)節(jié)點(diǎn)i插值，則在粗圖中找到距離節(jié)點(diǎn)i最近的k個(gè)節(jié)點(diǎn)，然后通過公式計(jì)算得到節(jié)點(diǎn)i的特征。最后，通過解碼器得到每個(gè)節(jié)點(diǎn)的速度與壓力信息。

02 論文源碼

原文的代碼總體上可以分為4個(gè)部分，分別是數(shù)據(jù)預(yù)處理、圖粗化層、圖卷積塊和圖恢復(fù)層。如圖3.2所示。

03 飛槳復(fù)現(xiàn)的技術(shù)要點(diǎn)

在模型中我們需要使用圖神經(jīng)網(wǎng)絡(luò)模擬網(wǎng)格，基于圖這種數(shù)據(jù)結(jié)構(gòu)模擬物理場(chǎng)的狀態(tài)，于是我們調(diào)用了飛槳的 PGL 庫。由于模型進(jìn)行圖數(shù)據(jù)的粗化時(shí)進(jìn)行了大量的稀疏矩陣運(yùn)算，為了提高運(yùn)算效率，減少運(yùn)算時(shí)間，我們需要對(duì)圖進(jìn)行合并。但是 PGL 是以列表的形式存儲(chǔ)圖，所以在圖數(shù)據(jù)進(jìn)入模型之前需要預(yù)處理，如下：

在基于代數(shù)多重網(wǎng)格的圖粗化層中，我們將圖粗化到不同的尺度，然后使用 GN 塊通過消息傳遞來總結(jié)和提取圖的特征。與基于代數(shù)多重網(wǎng)格的圖粗化層相比，圖恢復(fù)層采用反向操作。在圖恢復(fù)層中，我們使用空間插值方法對(duì)圖進(jìn)行上采樣，空間插值法的實(shí)現(xiàn)如下：

使用代數(shù)多重網(wǎng)格算法對(duì)圖進(jìn)行粗化時(shí)，涉及大量的稀疏矩陣乘法，于是我們調(diào)用了 scipy.sparse 庫。但是 scipy.sparse 庫中的.dot 方法和如今大部分深度學(xué)習(xí)框架中的稀疏矩陣乘法不同，后者會(huì)保留一些0，而前者只存非零值。這種差異會(huì)導(dǎo)致數(shù)據(jù)維度不一致，進(jìn)而影響之后的操作。針對(duì)這個(gè)問題，我們提出了兩種解決方案，第一種是采用填充的方法保持?jǐn)?shù)據(jù)維度，第二種是采用編碼器與解碼器將特征維度進(jìn)行壓縮。其中第一種方法與代數(shù)多重網(wǎng)格算法保持一致，精度較高。第二種方法提高了約20倍的訓(xùn)練速度，但是均方根誤差會(huì)略高。兩種方案的代碼如下：

• 方案1：

• 方案2:

預(yù)測(cè)結(jié)果

我們的模型按照論文給的案例，分別在機(jī)翼和圓柱這兩種不同的物理場(chǎng)景中進(jìn)行預(yù)測(cè)，其中機(jī)翼采用 NACA0012 翼型，表示翼型的網(wǎng)格包含6648個(gè)節(jié)點(diǎn)。圓柱體網(wǎng)格包含3887個(gè)節(jié)點(diǎn)。所有網(wǎng)格均由三角形和四邊形網(wǎng)格組成。翼型在穩(wěn)態(tài)、可壓縮和無粘性條件下由 Navier-Stokes 方程控制。圓柱在穩(wěn)態(tài)、不可壓縮和無粘性條件下由 Navier-Stokes 方程控制。預(yù)測(cè)結(jié)果表明，我們的模型預(yù)測(cè)的流場(chǎng)與真實(shí)流場(chǎng)基本相同，能夠達(dá)到原文的精度，證明復(fù)現(xiàn)成功。我們模型對(duì)圓柱和機(jī)翼的預(yù)測(cè)結(jié)果如圖4.1和圖4.2所示。原文的預(yù)測(cè)結(jié)果如圖4.3所示，其中原文的預(yù)測(cè)結(jié)果圖采用了專業(yè)的流體力學(xué)軟件繪制。

結(jié)語

這是我們團(tuán)隊(duì)第一次參加百度飛槳的論文復(fù)現(xiàn)賽，這次比賽的經(jīng)歷讓我們受益良多。首先，飛槳的復(fù)現(xiàn)賽賽事流程安排非常合理，我們被要求閱讀論文，然后跑通論文的原始參考代碼，這使得我們對(duì)論文的整體思路和原始代碼有了清晰的了解。在復(fù)現(xiàn)的過程中，我們快速了解兩種框架之間的差異，使用API映射表順利切換和使用不同的框架。盡管有一些 API 無法完全對(duì)應(yīng)，但我們可以進(jìn)行適當(dāng)?shù)馗膭?dòng)或重寫來解決這個(gè)問題。第二，飛槳團(tuán)隊(duì)為這個(gè)比賽提供了優(yōu)秀的平臺(tái)、充足的資源以及充分的指導(dǎo)，這使我們能夠更好地理解論文和實(shí)現(xiàn)代碼。他們的反饋對(duì)我們復(fù)現(xiàn)的結(jié)果和進(jìn)一步改進(jìn)非常有幫助。最后，特別感謝論文作者楊志雙學(xué)長(zhǎng)以及實(shí)驗(yàn)室的李天宇師兄對(duì)本次復(fù)現(xiàn)工作提供的技術(shù)指導(dǎo)和專業(yè)知識(shí)講解。

參與飛槳的論文復(fù)現(xiàn)賽是一次寶貴的學(xué)習(xí)和成長(zhǎng)機(jī)會(huì)。我學(xué)到了很多關(guān)于論文復(fù)現(xiàn)和深度學(xué)習(xí)框架轉(zhuǎn)換的知識(shí)，也結(jié)識(shí)了許多志同道合的同學(xué)和專業(yè)人士。未來，我們期待看到更多來自企業(yè)、高校、科研院所以及超算的小伙伴，加入飛槳 AI for Science 共創(chuàng)計(jì)劃，參與建設(shè)基于飛槳的 AI for Science 領(lǐng)域頂尖開源項(xiàng)目和開源社區(qū)。