華夏,柴志雷,2,張曦煌
(1.江南大學 人工智能與計算機學院,江蘇 無錫 214122;2.江蘇省模式識別與計算智能工程實驗室,江蘇 無錫 214122)
SNN(Spiking Neural Network)作為類腦計算的基礎,被譽為第三代人工神經網絡。追求卓越性能、探索適合SNN運行的類腦系統結構是SNN研究的主要任務之一。這意味著需要深入研究SNN模型及系統的工作負載特性,以設計高效的類腦計算平臺。由于SNN具有分布式計算特點,因此對SNN工作負載的研究需要合適的分布式計算平臺。在分布式平臺上進行SNN負載特性研究時,通常需要多次改變SNN模型規模,以及SNN工作負載與計算平臺之間的映射方式,以獲取不同情況下的工作負載數據。
因此,對分布式平臺進行內存負載分析能有效提升SNN工作負載研究的效率。一方面,對于特定SNN模型,通過分析滿足其內存需求的最小映射節點數,能夠幫助研究者確定計算平臺的可用映射方式的變化范圍。另一方面,通過分析計算平臺的網絡承載能力,幫助研究者在平臺網絡承載能力范圍內對其SNN模型進行網絡規模放縮,以研究不同規模下的SNN模型負載特性。
近年來,專用類腦系統和芯片不斷涌現,為SNN工作負載的研究提供了更多可選擇的平臺。2019年,施路平等人設計了Tianjic芯片,該芯片的采用了多核架構,能承載大約1 000萬個突觸模型和4萬個神經元模型。2020年,英特爾推出了包含768顆Loihi芯片的神經擬態系統Pohoiki Springs,該系統能夠承載約1億個神經元。
與專用類腦芯片一同發展的還有靈活性較強的SNN軟件仿真器。部分軟件仿真器構建的SNN模型能夠靈活地部署到硬件平臺中。文獻[6-8]研究了使用NEST等仿真器構建的SNN模型在類腦計算平臺上的工作負載特性。在這些研究中,使用者需要通過手動試探的方式獲取平臺的可用映射方式和網絡承載能力,這大大降低了SNN負載特性研究的效率。
針對現有的類腦計算平臺難以進行高效SNN工作負載研究的問題。本文首先對NEST仿真器進行工作負載特性分析,建立內存負載模型。其次搭建了PYNQ集群,并將NEST仿真器部署到集群中,實現了一種規模可擴展的分布式類腦計算平臺。最后設計了集群內存負載分析系統,通過提供內存消耗預測和集群網絡承載能力預測功能,幫助用戶建立SNN模型內存消耗的初步印象、快速確定集群網絡承載能力,以提高在類腦計算平臺上研究SNN工作負載的效率。
NEST是一款支持多種神經元、突觸模型的SNN仿真器,可以把大規模SNN模型映射到集群或者多核計算機上。
如圖1所示,NEST仿真器的運行流程分為三個階段。在創建和連接階段,NEST根據用戶指定的SNN網絡腳本構建神經元并建立神經元間的突觸連接。在仿真階段中,NEST循環執行突觸更新、神經元更新、脈沖傳遞等SNN模型仿真步驟。

圖1 NEST仿真器運行流程
NEST能夠進行分布式計算,在分布式計算平臺的每個節點中運行一個MPI進程,計算節點間使用MPI通信協議進行數據通信。使用NEST在分布式計算平臺上進行神經元構建和突觸構建時,分別遵循不同
的任務分配方式。
在創建階段中,NEST使用輪詢方式將神經元均勻分配到分布式平臺的不同節點上。NEST首先給每個神經元分發了一個全局標識編號(global identifier, GID),然后根據在分布式平臺上以循環方式分配神經元。以圖2為例,小型分布式計算平臺有兩個節點,對于共有4個神經元的SNN模型,按照NEST分配機制,平臺的每個節點分配2個神經元,IAF表示iaf_psc_alpha類型的神經元。

圖2 NEST神經元分配方式
突觸的分配與分布式平臺上的神經元分布有關,NEST按照突觸所對應的連接關系,將突觸分配到其連接目標所在的節點中。如圖3所示,在網絡連接階段,每個計算節點通過一張稀疏表來記錄突觸信息。在每個節點中,NEST需要遍歷節點分配到的每個神經元,然后對其中作為突觸的連接目標(突觸后)的神經元,根據突觸前后神經元GID在稀疏表對應位置記錄突觸信息。

圖3 NEST突觸分配方式
通過上述任務分配機制,NEST能夠有效利用分布式計算平臺的內存與計算資源,進行大規模SNN模型仿真。
根據1.1節分析,NEST運行時的內存負載消耗集中在建立階段和連接階段。這兩個階段中分別進行神經元和突觸的創建,其內存負載消耗主要包括神經元的內存占用M和突觸的內存占用M。此外,在分布式計算平臺上進行SNN模擬時,每個節點還需要為NEST提供初始化內存消耗M。因此,每個計算節點的總內存消耗由M、M和M組成,如式(1)所示:

其中,M包括NEST內核中數據結構初始化的內存消耗,其數值與計算平臺的性能有關。



表1 內存負載模型參數說明
每個計算節點的突觸內存占用M包括創建突觸模型及其基礎設施所需內存消耗,如式(3)所示。其中第一項為計算節點創建所有本地突觸所需的內存開銷,N(tgt)為突觸連接目標神經元tgt的數量,代表tgt的平均入度。第二項是計算節點為所有神經元提供的,用于記錄其連接關系的數據結構內存開銷。m、 分別是突觸模型及其基礎設施所需內存開銷,如表1所示。

為了提供SNN工作負載研究所需的分布式計算平臺。本節采用XilinxPYNQ-Z2開發板搭建集群,并在集群中部署了NEST仿真器,形成了一個規??蓴U展的分布式類腦計算平臺。
PYNQ-Z2開發板采用Zynq-7020作為其主芯片,該芯片通過數據接口將ARM Cortex-A9處理器核心(Processing System, PS)與可編程邏輯FPGA(Programmable Logic,PL)連接。
為了搭建分布式類腦計算平臺,本文首先設計了將NEST仿真器部署到單個PYNQ-Z2計算節點的方法。如圖4(a)所示,NEST仿真器的主體部分運行在PYNQ的PS端,用戶可以在PS端的JupyterNotebook中編輯SNN模型的腳本文件。Jupyter在運行該腳本文件的過程中,利用了PyNN庫提供的接口,通過SLI或者PyNEST界面完成對NEST內核中函數的調用。此外,利用PL端FPGA資源可以設計并行計算加速模塊,對NEST運行過程中的計算密集點進行加速運算,以提高NEST的運行速度。

圖4 PYNQ集群設計
如圖4(b)所示,PYNQ集群由三層硬件結構組成:PYNQ節點層、主控節點層和信息交互層。PYNQ節點層根據NEST任務分配機制處理SNN模型仿真過程中的計算任務,該層完全由PYNQ-Z2計算節點組成,每個節點都與信息交互層建立連接。信息交互層包含多臺交換機,負責傳遞計算節點間的信息。主控節點層由集群的主節點構成,主要負責將計算任務分配到每一個節點中。集群具有較強的規模擴展的能力,即計算節點規??梢噪S著計算需求的不同進行伸縮調整。
為了快速預測PYNQ集群運行SNN模型的內存消耗以及集群的網絡承載能力,本文設計了內存負載分析系統。
如圖5所示,內存負載分析系統主要包含五個模塊:SNN模型與平臺信息輸入、內存參數獲取模塊、SNN模型參數采集模塊、內存消耗預測模塊和網絡承載能力預測模塊。

圖5 內存負載分析系統結構
(1)SNN模型與平臺信息輸入:將描述SNN模型的腳本文件。
輸入系統,同時輸入計算節點總數、單節點最大內存M等計算平臺信息。
(2)內存參數獲取模塊:主要負責獲取表1所示的所有內存參數。
(3)SNN模型參數采集模塊:根據輸入的SNN腳本文件使用NEST構建該SNN網絡。在網絡構建(創建、連接階段)過程中,針對式(2)(3)中N(tgt)、(tgt)等與SNN模型相關的參數,調用統計功能進行采集。
(4)內存消耗預測模塊:根據內存負載模型,將獲取的所有內存參數、SNN模型參數轉化為內存消耗預測結果。此外,在滿足SNN模型內存需求的前提下,基于預測結果,給出計算平臺所需投入的最小節點規模。
(5)網絡承載能力預測模塊:根據內存參數和、M等平臺信息,預測計算平臺所能承載的最大SNN模型規模。
內存參數幾乎不受計算平臺和具體SNN模型的影響。采用分析與手動量化方法獲取所有內存參數。結果如表2所示。

表2 參數獲取結果
如圖6所示,SNN模型參數采集模塊內嵌在NEST仿真器的內核中,使用NEST構建SNN網絡時自動調用該模塊。輸入系統的SNN腳本文件通過接口nest.Create()和nest.Connect()調用NEST內核函數,進行神經元群落創建和突觸連接。同時這兩個接口也將各群落神經元規模,突觸連接規則等信息傳入NEST內核。SNN模型參數采集模塊使用2張參數表(群落表v_population和群落入度表v_in)記錄這些模型信息。每當調用nest.Create()創建神經元群落時,v_population記錄了此次創建的群落神經元數量。每當調用nest.Connect()創建突觸連接時,在v_in中記錄突觸連接目標神經元的入度(tgt)。

圖6 參數采集與預測
如圖6所示,內存消耗預測模塊首先讀取內存參數獲取、SNN模型參數采集模塊所獲得的全部參數數值,然后根據1.2節中的內存負載模型設計了負載計算函數,對所獲取的全部建模參數進行處理,最終預測出分布式計算平臺投入個節點運行SNN模型的總內存消耗()。
完成內存負載預測后,根據式(4)所示的穩定性能標準給出計算平臺為滿足模型內存需要所需投入的最小節點數,即最小性能穩定點。

其中M是計算平臺每個節點能夠提供的內存大小,是平臺能夠提供的最大節點數,所有滿足條件()≤M的節點規模均是可行節點規模,即能夠滿足網絡內存需求,保證平臺平穩運行。在所有可行節點規模中節點數量最小的即是最小性能穩定點Psteady。
內存消耗預測模塊中,預測的是SNN模型規模保持不變的情況下,平臺的內存消耗。而網絡承載能力預測,則是預測當計算平臺的可用內存均被占用時,所能承載SNN模型規模。
在分布式計算平臺上構建SNN模型時,若每個節點消耗的總內存都達到節點能提供的最大內存M,則SNN模型的規模達到了平臺網絡承載能力的極限。而單節點的總內存消耗包含M、M、M三部分,因此上述關系可以用公式(5)描述:

隨著投入計算節點數量的增加,分布式平臺的網絡承載能力也會提升。如公式(6)所示,根據1.2節建立的負載模型,可以將公式(5)轉化為平臺投入計算節點數與平臺網絡承載能力之間的關系,其中代表SNN模型的網絡規模,即模型中的神經元數量。

基于式(6),網絡承載能力預測模塊根據前面流程中獲取的內存參數、輸入系統的M等參數,預測投入個節點時的計算平臺能承載的最大SNN模型規模。
本文采用高性能計算機基準測試(HPC_benchmark)作為測試案例。該SNN模型中包含2個神經元群落,神經元群落間共有4組不同的連接關系,連接關系包含多條突觸。通過設定比例參數可以成倍調節模型網絡規模。本文選用的比例參數為0.5,此時網絡模型包含5 625個神經元和3 165萬條突觸。
為了驗證內存負載分析系統的有效性,本實驗針對HPC_benchmark,使用內存負載分析系統預測模型內存負載,并與實測數據進行對比。此外,利用系統的網絡承載能力預測功能,對分布式平臺在投入不同數量節點時的網絡承載能力進行預測,并且與實測平臺網絡承載能力進行了對比。
使用集群內存負載分析系統預測的內存消耗分布如圖7所示,由于HPC_benchmark中突觸的數量較多,因此突觸內存消耗M遠大于神經元內存消耗M。此外,隨著PYNQ集群投入的計算節點數增加,每個節點上的神經元、突觸內存消耗均被稀釋,而初始化內存消耗M的數值保持穩定。

圖7 內存消耗分布
圖8中,實線代表PYNQ集群投入不同數量節點運行HPC_benchmark時實測的內存消耗,虛線代表使用內存負載分析系統預測的數據。通過數據對比可知,使用內存負載分析系統預測的總內存消耗十分接近實測數據,平均而言系統的內存負載預測準確率達到了97.98%。此外,PYNQ集群單節點能提供最大內存約為450 MB,根據內存負載預測推測出的最小性能穩定點Psteady為3個節點,與實測數據相符。

圖8 內存消耗預測數據與實測數據對比
在本實驗中,首先通過手動反復調整HPC_benchmark網絡規模的方式,確定平臺投入不同數量節點時的真實網絡承載能力。然后將系統預測的平臺網絡承載能力與實測數據進行了對比。
圖9中,實線表示實際測得集群所能承載的網絡模型最大規模隨計算節點數的變化關系,虛線表示使用內存負載分析系統預測的數據隨計算節點數的變化關系。可用看出,預測的與實測數據十分接近。整體而言,內存負載分析系統在網絡承載能力預測方面的準確率約為97.19%。

圖9 網絡承載能力預測結果與實測數據對比
本文針對類腦計算平臺難以進行高效的SNN工作負載研究的問題,搭建了規模可伸縮的PYNQ集群類腦計算平臺,設計了內存負載分析系統,通過準確預測SNN模型在集群的上內存消耗以及集群的網絡承載能力,提高了在PYNQ集群分布式計算平臺上進行SNN工作負載的研究的效率。
在后續工作中,可以針對集群運行SNN模型的性能瓶頸進行優化,為SNN工作負載的研究提供高性能的平臺。