脈沖神經網絡(Spiking Neural Network, SNN)被譽為第三代神經網絡,使用更低層次的生物神經系統的抽象,既是神經科學中研究大腦原理的基本工具,又因其稀疏計算、事件驅動、超低功耗的特性而備受計算科學的關注。隨著深度學習方法的引入,SNN的性能得到大幅度提升,脈沖深度學習(Spiking Deep Learning)成為新興的研究熱點。傳統SNN框架更多的關注生物可解釋性,致力于構建精細脈沖神經元并仿真真實生物神經系統,并不支持自動微分,無法充分利用GPU的大規模并行計算能力,也缺乏對神經形態傳感器和計算芯片的支持。
為了解決上述問題,中國科學院自動化研究所李國齊研究員與北京大學計算機學院田永鴻教授團隊合作構建并開源了脈沖神經網絡深度學習框架SpikingJelly(中文名:驚蜇)。SpikingJelly(驚蜇)提供了全棧式的脈沖深度學習解決方案,支持神經形態數據處理、深度SNN的構建、替代梯度訓練、ANN轉換SNN、權重量化和神經形態芯片部署等功能。
SpikingJelly(驚蜇)框架具有以下幾個方面的主要優勢。
第一,簡單易用。脈沖深度學習作為計算神經科學與深度學習的交叉學科,要求研究者同時掌握兩個領域的知識體系,但實際情況是研究者可能只對一個領域有深度了解。SpikingJelly(驚蜇)框架提供了簡單易用的PyTorch風格的API、中英文雙語編寫的教程、活躍友善的討論社區,常用的網絡模型和訓練腳本也一并給出,研究者可以快速進行跨領域的學習和使用,通過寥寥數行代碼輕松構建并訓練深度SNN。
第二,擴展性強。通過模仿真實生物神經系統或借鑒人工神經網絡的成熟經驗,改進和提出新的SNN算法和模型,這種研究范式已經被廣泛使用,成果頗豐。研究者期望能夠自由地定義和擴展新模型,僅需少量代碼改動就對模型行為產生巨大改變,這一開發理念與SpikingJelly(驚蜇)的設計哲學不謀而合。SpikingJelly(驚蜇)中的絕大多數模塊都是通過層次清晰的多重繼承實現,這既為開發者降低了開發成本,也給用戶提供了完美的定義新模型的范例。
第三,性能卓越。深度學習涉及大規模數據處理、大尺度模型訓練,脈沖深度學習亦不例外。而SNN特有的額外的時間維度,使其具有更高的計算復雜度,因而脈沖深度學習對計算資源的需求更高。考慮到ImageNet這樣百萬規模樣本的數據集已經在SNN中廣泛使用,研究者愈發關注網絡的訓練速度。SpikingJelly(驚蜇)充分利用SNN的特性,通過計算圖遍歷順序優化、JIT (just-in-time compilation, 即時編譯)、半自動CUDA代碼生成等技術來加速SNN仿真,與其他框架相比可達11倍的訓練加速。第三方獨立社區Open Neuromorphic組織對多個SNN框架進行了精細的速度基準測試,參與評測的框架除了SpikingJelly(驚蜇)還包括來自Intel、SynSense時識科技、加州大學圣克魯斯分校、海德堡大學、瑞典皇家理工學院等科技公司和學術機構的框架,結果表明SpikingJelly(驚蜇)具有最快的仿真速度,比其他框架快10倍以上。
SpikingJelly(驚蜇)框架自2019年冬季一經推出就受到了研究者們的歡迎和廣泛使用,基于SpikingJelly(驚蜇)的研究工作已經大量出版,將SNN的應用從簡單的MNIST數據集分類擴展到人類水平的ImageNet圖像分類、網絡部署、事件相機數據處理等實際應用。此外,一些尖端前沿領域的探索也被報道,包括可校準的神經形態感知系統、神經形態憶阻器、事件驅動加速器硬件設計等。目前有超過95篇公開論文使用SpikingJelly(驚蜇)框架進行實驗,其中25篇發表在CCF-A類人工智能頂級會議,5篇發表在IEEE Transactions,1篇發表在Nature Communications。以上應用和研究表明,SpikingJelly(驚蜇)的開源,極大促進了脈沖深度學習領域的發展。
介紹SpikingJelly (驚蜇)框架的論文SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence于2023年10月6號在Science子刊《科學進展》(Science Advances)在線發表。北京大學計算機學院田永鴻教授、中國科學院自動化研究所李國齊研究員為論文共同通訊作者,北京大學計算機學院直博生方維為論文第一作者。國家自然科學基金、鵬城云腦網絡智能重大科技基礎設施項目、北京市自然科學基金為此項研究提供了資助。SpikingJelly(驚蜇)框架在2020-2022連續三年獲得科技部領導的啟智社區優秀開源項目,啟智社區亦為框架的發展提供了大力支持。
圖1. SpikingJelly論文發表于Science Advances
圖2. SpikingJelly(驚蜇)框架的整體結構、示例代碼、仿真速度、生態位以及典型應用
圖3. SpikingJelly(驚蜇)框架中的典型模塊
圖4:SpikingJelly(驚蜇)框架的典型應用
論文網址:
https://www.science.org/doi/full/10.1126/sciadv.adi1480
SpikingJelly(驚蜇)框架網址:
GitHub: https://github.com/fangwei123456/spikingjelly
啟智社區: https://openi.pcl.ac.cn/OpenI/spikingjelly
來源:中國科學院自動化研究所