脈沖神經網絡（Spiking Neural Network，SNN）被譽為第三代神經網絡，使用更低層次的生物神經系統(tǒng)的抽象。它既是神經科學中研究大腦原理的基本工具，又因稀疏計算、事件驅動、超低功耗的特性，備受計算科學的關注。隨著深度學習方法的引入，SNN的性能得到大幅提升，脈沖深度學習（Spiking Deep Learning）成為新興的研究熱點。傳統(tǒng)SNN框架更多地關注生物可解釋性，致力于構建精細脈沖神經元并仿真真實生物神經系統(tǒng)，不支持自動微分，無法充分利用GPU的大規(guī)模并行計算能力，并缺乏對神經形態(tài)傳感器和計算芯片的支持。?

中國科學院自動化研究所研究員李國齊與北京大學計算機學院教授田永鴻團隊合作，構建并開源了脈沖神經網絡深度學習框架SpikingJelly（中文名為“驚蜇”）?！绑@蜇”提供了全棧式的脈沖深度學習解決方案，支持了神經形態(tài)數(shù)據處理、深度SNN的構建、替代梯度訓練、ANN轉換SNN、權重量化和神經形態(tài)芯片部署等功能。?

“驚蜇”具有如下主要優(yōu)勢。第一，簡單易用。脈沖深度學習作為計算神經科學與深度學習的交叉學科，要求研究人員同時掌握兩個領域的知識體系，但實際情況是研究人員可能只對一個領域有深度了解。“驚蜇”提供了簡單易用的PyTorch風格的API、中英文雙語編寫的教程、活躍友善的討論社區(qū)，以及常用的網絡模型和訓練腳本，因此科研人員可以快速進行跨領域的學習和使用，通過寥寥數(shù)行代碼輕松構建并訓練深度SNN。第二，擴展性強。通過模仿真實生物神經系統(tǒng)或借鑒人工神經網絡的成熟經驗，改進和提出新的SNN算法和模型。這種研究范式已被廣泛使用、成果頗豐。研究人員期望能夠自由地定義和擴展新模型，僅需少量代碼改動便對模型行為產生改變。這一開發(fā)理念與“驚蜇”的設計哲學不謀而合?！绑@蜇”中的多數(shù)模塊均是通過層次清晰的多重繼承實現(xiàn)，既為開發(fā)者降低了開發(fā)成本，也給用戶提供了定義新模型的范例。第三，性能卓越。深度學習涉及大規(guī)模數(shù)據處理、大尺度模型訓練、脈沖深度學習。而SNN特有的額外的時間維度，使其具有更高的計算復雜度，因而脈沖深度學習對計算資源的需求更高?？紤]到ImageNet這樣百萬規(guī)模樣本的數(shù)據集已在SNN中廣泛使用，研究愈發(fā)關注網絡的訓練速度。“驚蜇”利用SNN的特性，通過計算圖遍歷順序優(yōu)化、JIT（just-in-time compilation，即時編譯）、半自動CUDA代碼生成等技術來加速SNN仿真，與其他框架相比可達11倍的訓練加速。第三方獨立社區(qū)Open Neuromorphic組織對多個SNN框架進行精細的速度基準測試。結果表明，“驚蜇”具有最快的仿真速度，比其他框架快10倍以上。

基于“驚蜇”的研究工作已大量出版，將SNN的應用從簡單的MNIST數(shù)據集分類擴展到人類水平的ImageNet圖像分類、網絡部署、事件相機數(shù)據處理等實際應用。此外，一些尖端前沿領域的探索亦被報道，包括可校準的神經形態(tài)感知系統(tǒng)、神經形態(tài)憶阻器、事件驅動加速器硬件設計等。上述應用和研究表明，“驚蜇”的開源促進了脈沖深度學習領域的發(fā)展。?

10月6日，關于“驚蜇”的研究論文（SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence），在線發(fā)表在《科學進展》（Science Advances）上。研究工作得到國家自然科學基金、“鵬城云腦”網絡智能重大科技基礎設施項目、北京市自然科學基金的支持。“驚蜇”在2020年至2022年連續(xù)獲得科技部領導的啟智社區(qū)優(yōu)秀開源項目。啟智社區(qū)亦為這一框架的發(fā)展提供了支持。

論文鏈接

“驚蜇”的整體結構、示例代碼、仿真速度、生態(tài)位及典型應用

“驚蜇”中的典型模塊

“驚蜇”的典型應用