SNNAX：基于 JAX 的脉冲神经网络
神经与进化计算
2025-01-28 v2 机器学习
摘要
脉冲神经网络（SNN）模拟器是构建生物启发模型和神经形态硬件架构原型并预测其性能的重要工具。对于此类工具而言，易用性和灵活性至关重要，但模拟速度同样关键，尤其是在 SNN 模拟固有的复杂性下。在此，我们提出 SNNAX，一个基于 JAX 的框架，以具备类似 PyTorch 的直观性和类似 JAX 的执行速度来模拟和训练此类模型。SNNAX 模型易于扩展和定制，以适应所需的模型规格和目标神经形态硬件。此外，SNNAX 提供了优化 SNN 训练和部署的关键特性，如灵活的自动微分和即时编译。我们评估并将 SNNAX 与其他常用于编程 SNN 的机器学习（ML）框架进行比较。我们提供了在 SNNAX 中模拟 SNN 的关键性能指标、最佳实践、文档化示例，并实现了文献中使用的多个基准测试。
引用
@article{arxiv.2409.02842,
title = {SNNAX -- Spiking Neural Networks in JAX},
author = {Jamie Lohoff and Jan Finkbeiner and Emre Neftci},
journal= {arXiv preprint arXiv:2409.02842},
year = {2025}
}