開源脈衝神經網路深度學習框架——驚蟄(SpikingJelly)
開源脈衝神經網路深度學習框架——驚蟄(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將常用神經形態資料集統一進行了封裝，只需一行程式碼即可進行呼叫，其中事件資料被統一為(