基于SNN 神經元重分布的NEST 仿真器性能優化

劉家航，郁龔健，李佩琦，華 夏，柴志雷，2，陳聞杰

（1.江南大學 人工智能與計算機學院，江蘇 無錫 214122；2.江蘇省模式識別與計算智能工程實驗室，江蘇 無錫 214122；3.華東師范大學軟硬件協同設計技術與應用教育部工程研究中心，上海 200062）

0 概述

深度學習在數據挖掘、自然語言處理、語音處理、文本分類等諸多領域都取得了突出的成果，但其存在需要海量標注數據、通用智能水平弱等局限性［1］，難以成為解決人工智能問題的終極手段。人類大腦是由約1011個神經元、1015個突觸構成的復雜的生物體，具有很高的智能水平，但功耗只有25 W左右［2］，其計算模式值得借鑒。目前，出現了越來越多的大腦神經計算科學的研究工作，以期通過解析人類大腦工作機理、發展類腦計算來克服現有深度學習的不足［3］。類腦計算的基礎是脈沖神經網絡（Spiking Neural Network，SNN），與傳統的深度神經網絡（Deep Neural Network，DNN）相比，SNN 的工作機理更接近于生物大腦。

大腦規模所需的計算資源遠超單個計算節點或芯片的功能［4］，因此，為了獲得最佳的計算能力，構建大規模集群形成分布式計算平臺成為當前主流［5-7］。然而隨著計算節點的增多，通信時間所占比例越來越大，有效計算的比例逐漸下降，導致SNN的計算效率下降。

針對通信的優化，壓縮、打包尖峰數據等［8-10］傳統方法只關注對通信的脈沖數據進行優化，局限性較大。在神經元映射對通信的影響方面，CARLOS等提出HP-NBX 方法［11］，利用超圖分區和動態稀疏交換減少進程之間的連接，從而提高通信平衡性，優化通信過程。URGESE 等提出SNN-PP方法［12］，在SpiNNaker 中使用譜聚類算法將神經元分組成亞群，使亞群中緊密相連的組保持在同一個計算節點（進程）中，目的是最小化相互的連接并保持高度連接的組件相互靠近。BALAJI 等提出SpiNeCluster 方法［13］，根據尖峰信息將SNN 劃分為局部和全局突觸，最大限度減少共享互連上的尖峰數量。XIANG等提出一種基于跨層的神經映射方法［14］，將屬于相鄰層的突觸連接神經元映射到同一片上網絡節點，同時為適應各種輸入模式，還考慮輸入尖峰率并重新映射神經元以提高映射效率。

針對大規模分布式類腦計算的仿真，NEST 仿真器是一款主流的類腦模擬器開源軟件［15］。NEST 仿真器可以模擬任何規模的脈沖神經網絡，并且支持集成式的MPI、OpenMP 通信協議，其通過分布式計算大幅提高了仿真速度。NEST 的通信機制采用“緩沖區動態相等”的方法，各節點無需通信，直接動態保持發送緩沖區相等。

雖然NEST 對通信時間進行了一定程度的優化，但是由于緩沖區互相無交流，使得通信數據量持續增加，因此其在能耗方面表現較差。針對這一問題，本文提出基于SNN 子圖跨節點優化的神經元重分布算法ReLOC，通過分析SNN 模型中神經元之間聯系的緊密性，減少神經元到其他節點的連接，并以稀疏交換思想對NEST 本身的通信機制進行改進，利用其稀疏性提高通信效率。

1 研究背景與相關工作

1.1 脈沖神經網絡

脈沖神經網絡（SNN）被譽為第三代人工神經網絡［16］，其模擬神經元更加接近實際，并且把時間的影響也考慮在內。脈沖神經網絡中的神經元不是在每一次迭代中都被激活，而是在其膜電位達到某一個特定值時才被激活。當一個神經元被激活時，它會產生一個信號傳遞給其他神經元，并改變其膜電位。

神經科學家通常用SNN 評估理論和實驗結果［17-18］。SNN 仿真分為2 個階段：計算和交換。計算階段包括基于偏微分方程更新神經元模型；交換階段將神經元產生的脈沖發送到其突觸后神經元處。在分布式系統中，交換涉及到進程間通信過程，而節點規模越大，通信所占的比例就越大。

1.2 NEST 仿真器

NEST 是一款應用廣泛的SNN 模型仿真器，其大規模分布式計算特性非常適合用來作為SNN 負載表征的模擬器工具。本節介紹NEST 的仿真架構、并行分配和通信機制。

1.2.1 NEST 仿真架構

圖1 展示了NEST 的運行過程。首先使用Create在各個節點上創建神經元；然后Connect 連接神經元，創建突觸；最后Simulate 進行仿真，此過程分為三步，即更新所有神經元變量（Update）、將產生的脈沖發送出去（Communicate）以及更新本地突觸變量（Deliver）。

圖1 NEST 運行流程Fig.1 Procedure of NEST operation

1.2.2 NEST 并行架構和神經元分配

NEST 并行架構采用支持MPI 節點間通信和OPENMP 線程通信，并使用虛擬進程（Virtual Process，VP）的概念來統一管理并行系統，虛擬進程數量NVP由進程數量P和線程數量T決定，NVP=P×T。

神經元分配采用基于VP 的循環分配，使得神經元在各個節點上均勻分布。如圖2 所示，其中：Sg 代表spike-generator 脈沖激發器，此設備在VP 上是完整創建的；Iaf 代表神經元，是基于虛擬進程的分配。由于這種分配方式不可避免地將聯系緊密的神經元分配到不同的節點，導致通信消耗增大，因此神經元分配對于通信的改進是必要的。

圖2 NEST 并行架構和神經元分配方式Fig.2 Parallel architecture and neuron allocation mode of NEST

1.2.3 NEST 通信架構

NEST 采用將單個時間片內的神經元連續更新的思想：神經元先進行若干次連續更新，然后一次時間片內產生的脈沖統一進行MPI 通信。

NEST 使用MPI_Alltoall 函數進行通信，通常的設計是每次通信前先通信一次確保每個節點的發送緩沖區都保持相等，這樣通信兩次的設計毫無疑問會隨著節點增加造成通信時間劇增，因此，NEST 不使用這種設計，而是采用所有節點同時自主保持通信緩沖區相等的策略。

如圖3 所示，在進行每一輪通信時，首先將脈沖放入發送緩沖區中，如果緩沖區的容量小于脈沖數，則將脈沖放滿進行通信，然后擴大發送緩沖區的容量，繼續放入剩下的脈沖，直到脈沖傳輸完成，其中發送緩沖區的容量會一直變大。這樣設計提高了通信的可擴展性，但是會導致通信能耗迅速增加，并且NEST 集體通信導致發送緩沖區無論有沒有脈沖都需要發送數據，這無疑增加了通信消耗，而并行SNN 的通信應當是進程間存在脈沖交換時才會發生。

圖3 NEST 通信機制流程Fig.3 Procedure of NEST communication mechanism

1.2.4 NEST 負載特性

在使用多節點仿真時，神經元會被均勻地分發到每個節點上，產生跨節點通信。如圖4 所示，在NEST 上運行平衡隨機網絡模型，隨著節點增加，通信所占比例逐漸增大，計算效率逐漸下降。

圖4 NEST 各階段時間占比Fig.4 Time proportion of NEST each stage

圖5 顯示了NEST 通信能耗的變化，可以看出，隨著節點的增加，通信數據量迅速上升，NEST 在通信能耗方面可擴展性變差。針對通信能耗可擴展差的問題，本文對NEST 進行稀疏性通信優化，以期提高通信效率。

圖5 通信數據量Fig.5 Communication data volume

2 稀疏性通信優化設計與實現

針對NEST 集群的神經元循環分配帶來的問題，本文實現神經元重分布算法ReLOC。通過對神經元之間連接關系的分析，重新分布神經元，使得跨硬件節點之間的連接變得更加稀疏。在此基礎上，改進NEST 通信方式，使NEST 集群可以更好地利用稀疏性以減少通信代價。

2.1 影響通信的因素

分布式集群最重要的問題就是節點間的通信問題，因為突觸前和突觸后神經元總是不可避免地被分布到不同的節點，最優分布是所有神經元都被放入同一個節點，但是此分布無法體現分布式集群的優勢。節點間通常使用以太網連接，因此，網絡帶寬和延遲直接影響通信速度。除此之外，神經元彼此之間的連接也會對跨節點通信造成較大影響。

將一個SNN 模型建立為如圖6 所示的SNN 連通圖，圓點代表神經元，中間帶箭頭的線段表示突觸。然后將一個神經元和它的突觸后神經元共同組成這個SNN 圖的子圖，如圖7 所示，假如一個神經元產生脈沖，那么脈沖數據會被發送給所有的突觸后神經元。所有的脈沖傳輸都是基于這樣的SNN子圖。

圖6 并行SNN 連通圖Fig.6 Parallel SNN connected graph

圖7 SNN 子圖Fig.7 SNN subgraph

在NEST 仿真器中，當一個節點上包含一個神經元的多個突觸后神經元時，這些突觸后神經元會共享一個脈沖數據。因此，無論突觸后神經元的數量如何，此神經元發出的脈沖傳輸到此節點都只算作一個，即神經元到節點的連接。每次神經元更新產生的總脈沖總是小于或等于節點數乘以神經元數。因此，如何將一個SNN 子圖中的神經元分布到盡可能少的節點上，并且改進NEST 通信機制使其更適用稀疏性優化成為重要問題，本文設計的ReLOC 算法每次都加入當前匹配集最優匹配的神經元，形成新的匹配集。根據最優匹配規則選取屬于匹配集中的突觸后神經元最多的神經元，將SNN子圖放入盡可能少的節點，從而減少跨節點脈沖。在此基礎上，改進NEST 原本的集體通信，利用稀疏交換的思想，在跨節點脈沖減少的同時，避免多余的進程數據交換，從而提高NEST 通信效率。

2.2 ReLOC 算法

2.2.1 算法描述

ReLOC 算法描述如算法1 所示，其將神經元按照緊密度的聯系分布到若干節點上，通過對節點神經元容納量的限制，平衡節點的負載。

算法1ReLOC 算法

ReLOC 算法輸入為節點數和神經元連接表。一個SNN 模型可以建立成一張連通圖G=（V，E），V代表所有的神經元集合，E代表神經元之間的突觸連接。而神經元連接表可以表示為圖G的鄰接矩陣A，A是Nn×Nn維的，Nn表示神經元的總數，節點數用Nnp來表示。ReLOC 算法輸出為各節點的神經元分布，表示為一個分布矩陣T，T是Nn×Nnp維的。

首先選取一個節點，將一個隨機的神經元r分布到此節點上，k表示當前節點；然后依次在此節點上放入未分配的神經元，直到節點k達到最大容納量max_neurons，規則是每次放入的神經元都是未分配的神經元集neurons 中神經元的突觸后神經元在節點k中分布最多的神經元。當神經元按照此方法依次進行匹配時，每次匹配的神經元都是與匹配集聯系最密切的神經元。ReLOC 算法中的符號定義如表1 所示。

表1 符號定義Table 1 Symbolic interpretation

ReLOC 算法的思想是先選定一個神經元，每次匹配一個和它聯系最緊密的神經元，形成一個神經元集，然后繼續循環匹配和此神經元集聯系最緊密的神經元，生成新的神經元集。此算法中決定神經元之間密切程度的是神經元集中屬于突觸后神經元的數量。因為每次匹配都會使一個SNN 子圖（見圖7）最大化分布到同一個節點上，當所有匹配都是最優時，每個子圖都會避免分布在多個節點上，從而減少通信量。

ReLOC 算法只需要在任意一個節點上運行，然后將分布結果發送給其他節點即可。內存使用按照式（1）計算：

其中：M表示總內存；N表示神經元總數量；Mint表示整型參數占用內存。

2.2.2 ReLOC 算法的有效性

ReLOC 算法的核心是匹配過程，每一次匹配都會選擇最優的匹配集，最優意味著最低代價的SNN子圖切割，即SNN 子圖分布在最少的節點上。

為了描述SNN 的拓撲，將神經元連接圖用鄰接矩陣A來表示，A[i，j]∈{0，1}，當神經元i和j之間存在突觸連接時，A[i，j]=1，否則A[i，j]=0。用T來表示分布矩陣，當神經元i被分布到節點j時，T[i，j]=1，否則T[i，j]=0。

根據上文提到的突觸后神經元共享脈沖數據，衡量跨節點通信可以用神經元到節點的連接數量NNP來表示。首先定義P：

P[i，j]≥0 表示神經元i到節點j存在連接，那么NNP可以表示為：

2.2.3 基于ReLOC 算法的NEST 性能優化

為了利用ReLOC 算法的便捷性，并且與原先的NEST 調用代碼不產生沖突，本文改變了NEST 的基本架構，在Connect 時不創建具體的突觸對象，只是通過所有連接得到一個連接表，然后在調用Simulate函數初始化仿真時根據連接表進行神經元的重分布處理，之后再創建突觸和運行仿真過程。重分布具體流程如圖8 所示。其中，Setnod 是系統調用ReLOC 進行神經元重新分配的標志，對于非稀疏性SNN 模型，則無需調用此模塊。Setnod 函數可以輸入參數，根據參數選擇不同的重分布算法，此模塊是可擴展的，可以自行添加新的算法。

圖8 重分布流程Fig.8 Procedure of relocation

用戶調用應用重分布稀疏方法的NEST 內核的Python 代碼以如下為例，對比原代碼新增了2 個函數調用，設置重分布（Setnod）和設置通信方式（Setcomm），兩者都需要在ResetKernel 后進行調用。如果不調用則默認使用原來的NEST 神經元分配和通信方式。

2.3 稀疏交換

稀疏交換是HOEFLER 等提出的動態稀疏數據交換算法［19］，是節點對節點通信的模式，即一個節點只發送必要的數據給另一個節點，期間需要啟動多次MPI通信，總體效率不如集體通信（MPI_Alltoall）。但是如果減少每個節點的通信節點（一個節點需要向另一個節點發送脈沖數據，即本節點神經元到其他節點的連接）的數量，那么就可以減少啟動MPI 通信的次數，進而提高通信效率。

如圖9 所示，在啟動通信階段，每個節點都會使用MPI_Reduce_scatter 函數與其他節點交流得到源節點的信息，避免與非相鄰進程交流，以此提高通信效率。然后每個節點都將脈沖數據異步發送給其所有的目標節點，并等待接收其源節點發送的脈沖，一旦接收完畢，則表明通信階段結束。

圖9 稀疏交換流程Fig.9 Procedure of sparse exchange

每個節點與其他目標節點進行數據傳輸時使用非阻塞通信，即用MPI_Isend 進行發送數據，用MPI_Recv 接收數據，最后使用MPI_Waitall 等待數據全部傳輸完成，然后繼續仿真操作。此過程不需要等待全部節點通信結束，通過非阻塞通信、重疊通信與計算階段提高并行計算的效率。

在NEST 上使用稀疏交換之前需要先調用Setcomm（type）函數，其中的參數可以設置，從而啟用不同的通信機制，如果不調用此函數則默認通信為NEST 原本通信模式。

在使用稀疏交換之前先進行神經元重分布操作，減少本節點神經元到其他節點的連接，由于稀疏SNN 模型每一輪神經元產生的脈沖比較稀少，因此減少了節點到節點通信的發生，并且只需要發送目標進程所需的脈沖數據，一定程度上解決了NEST集體通信帶來的通信能耗可擴展性差的問題。

3 實驗結果與分析

3.1 實驗案例選擇

本文選擇以下2 個實驗案例進行分析：

1）皮質微電路SNN 模型

此案例基于POTJANS 等提出的皮質微電路（Cortical Microcircuit，CM）模型［20］。CM 模型網絡包含4層，每層由抑制性和興奮性神經元群組成，分為8個群體，共7.7 萬個神經元和3 億突觸。神經元類型為iaf_psc_exp，所有連接均為靜態突觸，連接規則為fixed_total_number，另有8 個泊松和對應的神經元群全連接。此案例在生物學上比較合理地模擬了大腦皮層的神經活動，為了驗證網絡稀疏性對于SNN 分布式計算的好處，將該模型的神經元和突觸縮放至0.1和0.02，以0.1 ms 的時間步長仿真200 ms。

2）平衡隨機網絡模型

在此案例中，平衡隨機網絡模型的神經元被分成2 個種群，分別由400 個興奮性神經元和100 個抑制性神經元構成，每個神經元的固定入度為5，因此，共有2 500 個突觸。神經元類型為iaf_psc_alpha，所有連接均為靜態突觸，連接規則為fixed_indegree，1 個泊松和所有神經元連接。將該模型分別在類腦平臺上以0.1 ms 的時間步長仿真200 ms。

3.2 實驗平臺

本文選取由28 塊PYNQ-Z2 組成的高性能異構類腦平臺，每個節點包括PS（Process System）端的ARM A9 雙核處理器系統和1 個PL（可編程邏輯）端的FPGA 器件，物理內存512 MB。節點之間使用1 000 Mb/s 網絡帶寬的以太網進行通信，采用TCP/IP 協議。

3.3 實驗內容

上述2 種模型都是Python 編寫的NEST 案例，本文調用NEST 的內核（C++實現）對SNN 進行仿真，主要特征是采用時間步驅動，并行通信使用MPI 庫。

對模型進行時間統計：1）計算時間是指每個節點所有神經元順序狀態更新所花時間；2）通信時間是指節點通信啟動到數據接收完成所花的時間；3）仿真時間是指完成仿真所花的全部時間，包括計算時間和通信時間；4）通信數據量是指MPI 函數在200 ms 仿真時間內所有節點實際發送的全部數據字節量（unsigned int 型）。

3.4 結果與分析

3.4.1 神經元重分布

分別測試應用神經元重分布算法和循環分布算法的交換數據量和神經元到進程的數量，通信模式均采用NEST 默認通信機制。如圖10（a）所示，在應用ReLOC 算法后，神經元到進程的數量平均減少約20%，但是由于NEST 的通信機制為集體通信，無法利用稀疏連接帶來的好處，因此交換的數據量并未減少，如圖10（b）所示。

圖10 重分布和循環分布算法實驗結果對比Fig.10 Comparison of experimental results by relocation and cyclic allocation algorithms

3.4.2 神經元重分布和稀疏交換

使用重分布稀疏交換方法（ReL-SE）和原版本NEST 模式運行3.1 節的2 個實驗案例，這2 個案例對于僅使用靜態突觸的稀疏性SNN 網絡有比較好的代表性。分別使用4、8、16、28 個節點進行仿真，運行結果如圖11 和圖12 所示。可以發現，對于稀疏SNN 模型，在通信數據量大幅減少的情況下，通信時間和仿真時間與原版本相差不大，這就使得通信能耗大幅減少。如圖11 所示，在皮質微電路模型中，由于起初神經元脈沖比較密集，導致發送緩沖區一開始就會非常大，因此后續通信數據的浪費更加明顯，在28 個節點上的仿真結果表明，使用重分布稀疏交換方法的通信能耗相比NEST 模式減少約98.63%。由圖11（a）和圖12（a）可以看出，當節點較少時，NEST 模式的仿真時間少于重分布稀疏交換方法，然而節點繼續增加，NEST 模式的仿真時間卻呈快速上升趨勢。與此相應，通信時間也在快速接近，可以預見的是，節點的擴張有利于稀疏通信。由圖11（c）和圖12（c）可以看出，NEST 原本通信數據量隨著節點的增加呈現快速上升趨勢，而稀疏交換的數據量卻緩慢增長，由此可見，同時配后稀疏交換能夠大幅降低通信能耗。

圖11 皮質微電路模型運行結果Fig.11 Running result of CM model

圖12 平衡隨機網絡模型運行結果Fig.12 Running result of balanced stochastic network model

3.5 ReLOC 算法的實現代價

如圖13 所示，隨著節點的增加，ReLOC 算法花費時間越來越少，這是因為節點的增加導致每個節點上所能容納的神經元減少，因此在算法循環進行神經元匹配時，突觸后神經元位于節點神經元集的個數的計算次數大幅減少。

圖13 ReLOC 算法耗時Fig.13 Time consuming of ReLOC algorithm

由于神經元重分布是仿真前一次性消耗的時間，因此隨著仿真時間的增加，ReLOC 所需代價逐漸降低。所以，節點和仿真時間與ReLOC 算法的代價成反比，使得代價最終可以忽略不計。

4 結束語

NEST 循環分配導致聯系緊密的神經元易產生分離。針對該問題，本文通過實現神經元的重分布，使聯系緊密的神經元更容易分布到同一節點，使進程間的通信更加稀疏。同時，對NEST 本身的通信機制進行改進，利用稀疏交換思想大幅減少通信數據量，并配合ReLOC 算法提高通信效率。本文在使用重分布稀疏交換方法時只需要調用2 個函數來啟動使用，與原版本NEST 的用戶交互相比并未發生很大的變化，此方法的效果隨著節點增多進一步明顯，而實現代價隨節點和仿真時間的增加逐漸減少。后續將通過對神經元進行更有效的分配，進一步縮短通信時間。同時，由于本文算法對于擁有STDP 突觸的案例效果不佳，因此下一步也將考慮2 種突觸的平衡性優化神經元重分布算法。