Feb, 2024

Spyx:用于脉冲神经网络的即时编译优化的库

TL;DR人工智能在现代社会中的作用日益重要,高效训练和部署深度神经网络成为关注的重点之一。脉冲神经网络(SNNs)通过利用时间稀疏计算提供了提高能源效率的方法,但其训练由于其循环性质而具有一定挑战性。本文介绍了 Spyx,这是一个基于 JAX 设计的新型轻量级 SNN 模拟和优化库,通过在现代加速器的扩展 vRAM 中进行预存数据,并采用大量的即时编译,Spyx 使得 SNN 优化能够作为统一的低级程序在 NVIDIA GPU 或 Google TPU 上执行,实现了最佳的硬件利用率,并在保持相当灵活性的同时超越了许多现有的 SNN 训练框架的性能。