開源脈沖神經網絡深度學習框架——驚蟄(SpikingJelly)

背景

　　近年來神經形態計算芯片發展迅速，大量高校企業團隊跟進，這樣的芯片運行SNN的能效比與速度都超越了傳統的通用計算設備。相應的，神經形態感知芯片也發展迅速。目前已有各種模態的感知芯片，其中如北京大學黃鐵軍教授團隊的Vidar相機，功能上仿照視網膜中央凹，能輸出脈沖信號，高速情況下實現比傳統相機更清晰的采樣。脈沖網絡研究領域頂會文章與Nature Science刊物文章也在逐年增長(如下圖)。通過ANN轉換SNN，SNN首次達到媲美ANN的性能。同時，隨着梯度替代(surrogate gradient)法的提出，直接利用梯度下降法進行SNN訓練成為可能。目前，利用SNN進行深度學習已經成為機器學習領域的研究熱點，在視覺分類，目標檢測，強化學習等領域取得了不錯的成果。

 

*僅計算Research與Reviews文章

 

　　SpikingJelly的前身SpikingFlow，作為北京大學本科生《神經網絡的計算基礎》& 研究生《機器學習原理》課程教學實驗平台，已經有將近一年的時間。SpikingFlow與SpikingJelly的關系猶如Minix之於Linux，前者作為教學用的示例，后者在前者的基礎上改進，吸納教學過程中的意見，新的工程實踐和研究成果。

 

　　如上圖所示，目前已有的機器學習框架和脈沖網絡仿真框架都無法彼此兼容，傳統的脈沖網絡框架無法與當下成熟的深度學習技術相結合。因此，SpikingJelly也就應運而生。同時期，國外也誕生了類似的Norse框架，兩者達成了良性互動，共同推進了該領域的發展。

圖源: Deep learning incorporating biologically inspired neural dynamics and in-memory computing

 

脈沖神經網絡介紹

　　脈沖神經網絡，簡稱SNN，被譽為第三代人工神經網絡，是由大腦這樣一個脈沖信號處理系統啟發而構建的。大腦具有高級智能，並且功耗較低(有一種說法是僅相當於一個25w的燈泡)。通過借鑒大腦中的脈沖結構，SNN能夠在保持低功耗的前提下，達到與ANN相當的性能。

　　SNN的結構仿照了生物神經系統的組織結構，現在使用的深度SNN通常采用前向結構(如下圖所示)。相比前向ANN，SNN的復雜性不僅體現在連接權重與拓撲的多樣性，還體現在神經元內在的動力學方程上，能夠同時處理時間域與空間域的信息。

　　

　　基於SNN的諸多特性，其應用前景十分廣闊，從簡單的圖像分類，動作識別，音頻處理到復雜的音視頻信號，強化學習，機器人控制任務，都是SNN施展自己本領的潛在舞台。

圖源: TrueNorth: Accelerating From Zero to 64 Million Neurons in 10 Years

 

框架概況

　　驚蟄(SpikingJelly)是一個開源脈沖神經網絡深度學習框架(框架主頁：https://github.com/fangwei123456/spikingjelly (GitHub)；https://git.openi.org.cn/OpenI/spikingjelly (OpenI))。SpikingJelly框架整體結構圖如下，使用PyTorch作為自動微分后端，利用C++和CUDA擴展進行性能增強，同時支持CPU和GPU。框架中包含數據集，可視化，深度學習三大模塊。目前社區主要由北大媒體學習組和鵬城實驗室人工智能中心運營管理。

 

 　　在SpikingJelly框架中，神經元的動態被描述為充電，放電，重置三個過程，與圖中三種顏色分別對應。

 

 　　框架支持梯度替代法與ANN轉SNN法兩種主流SNN的深度學習算法，也是在實際任務上目前性能最好的兩種算法。之前介紹的神經元動態方程中，只有放電過程是不可微分的，為了解決這個問題，近年來領域內提出了替代梯度方法直接訓練SNN，原理是梯度在反向傳播時采用替代函數來近似放電過程中的階躍函數。另一種深度學習方法，從ANN轉換為相應的SNN，也受到了重點關注。這是由於ANN中的ReLU神經元非線性激活和SNN中IF神經元的發放率有着極強的相關性，可以基於類似的相關性將訓練好的ANN轉換為對應的SNN，省去了直接訓練SNN的難題。

　　框架提供了Neuron，Layer，Functional，Encoding四個基本模塊：

• Neuron提供了深度SNN中最常用的LIF和IF神經元；
• Layer提供了SNN中特有的網絡層；
• Functional提供深度SNN所需的函數；
• Encoding提供常用的脈沖編碼器。

　　框架還集成了多個神經形態數據集，因為大多數神經形態數據集需要專用軟件讀取，使用非常繁瑣。因此SpikingJelly將常用神經形態數據集統一進行了封裝，只需一行代碼即可進行調用，其中事件數據被統一為(𝑥,𝑦,𝑡,𝑝)的Python字典格式，並提供將事件積分成幀數據的接口。來自其它框架，其他類型的數據也可以通過編碼器間接使用。為了方便下載，在遵守原始協議或原作者同意的情況下，在OpenI上建立了數據集的下載鏡像。

　　框架還提供了一些對神經元狀態，脈沖特征進行可視化的API，方便研究者直接使用，不需要再去進行繁瑣的繪圖調整。

 

　　框架還提供了豐富的教程，通俗易懂，初學者和資深開發者均可快速上手，包含了方便初學者上手學習的基礎教程，和實現具體任務的樣例教程。

　　SpikingJelly的模塊可以無縫嵌入到PyTorch的ANN模塊中，學習成本很低，只要使用者事先掌握PyTorch，就能夠輕松上手SpikingJelly，框架還提供了詳細的雙語文檔，部分經典算法還提供了原始論文的相關內容，方便用戶在開發的同時學習領域內的研究成果。

　　由於SNN的發放脈沖特征圖中絕大部分位置都是0，針對SNN發放脈沖稀疏的特點，框架中加入了一些初步的工作，其中主要是稀疏矩陣運算的加速。根據實測，在矩陣規模大，且稀疏度高時，稀疏矩陣乘法的速度遠高於普通矩陣乘法。

　　此外，SpikingJelly的模塊在任務性能上也高度優化，不需要復雜的設計，只需要簡單的網絡和基本的神經元就可以達到STOA的性能。

 

　　在常用的benchmark上，框架實測ANN轉SNN的性能損失很低，這使得很多不了解SNN的人也可以將他們的工作轉換為相應的SNN，並達到接近的性能。

 

代碼風格 

　　在神經元的實現中，隱狀態被封裝在模塊內部，而PyTorch, Norse等框架的狀態定義模塊外部，這可以避免使用者手動維護神經元膜電位，並且在搭建多層網絡時可以直接堆疊。

　　

　　前面介紹的直接堆疊就是這里的逐步傳播方式，支持用Sequential直接包裝網絡，使用簡單。框架同樣也提供了逐層傳播的方式，並行度更高，但也會帶來稍大的顯存開銷。兩種方法在前向，反向傳播時均完全等價。

　　下面是一個簡單的網絡示例，熟悉ANN的用戶能夠很快上手SNN的搭建。

 

　　ANN2SNN模塊支持兩種后端模型表示，同時也支持用戶自定義對網絡進行轉換的方式。