mlGeNN: accelerating SNN inference using GPU-enabled neural networks
郑重声明：原文参见标题，如有侵权，请联系作者，将会撤销发布！
Neuromorphic Computing and Engineering, 2022
Abstract
在本文中，我们介绍了mlGeNN —— 一个Python库，用于将Keras中指定的人工神经网络(ANN)转换为脉冲神经网络(SNN)。SNN使用带有扩展的GeNN进行模拟，以有效地支持卷积连接和批处理。我们在CIFAR-10和ImageNet分类任务上评估转换后的SNN，并将性能与原始ANN和其他SNN模拟器进行比较。我们发现使用在CIFAR-10数据集上训练的VGG-16模型执行推理比BindsNet快2.5，当使用在CIFAR-10上训练的ResNet-20模型和FewSpike ANN到SNN转换时，mlGeNN仅比TensorFlow慢2倍多一点。
1. Introduction
对于图像分类，在ImageNet2012 [1]等大型数据集上训练的深度卷积神经网络(CNN)是黄金标准。然而，即使在训练CNN之后，使用它进行推理也需要数百万(如果不是数十亿)的乘法累加操作(参见Hanif等人[2]，图13.5)通过多层训练的权重来传播实值单元激活。这导致计算和能源[3]要求使其不适用于许多移动和边缘计算应用程序，并引发了对神经网络算法碳足迹的担忧。已采用各种方法来减少计算时间和能源成本，包括简化模型结构[4]、使用低精度数据类型[5]甚至专用加速器硬件[6]。
脉冲神经网络(SNN)代表了一种更根本的方法变化。从生物大脑中获得更接近的灵感，SNN中的神经元使用被称为"脉冲"的离散事件的时空模式进行通信，而不是深度网络中单元交换的实值"发放率"。因为通常在任何给定时间，SNN中只有一小部分神经元处于活动状态(脉冲)，所以为单个时间步骤模拟SNN模型需要的操作要少得多。此外，由于每个脉冲神经元的激活都是二值的，因此通过一层权重传播它只需要加法而不是乘法累加操作，这可能允许更高效的硬件实现。已经开发了各种各样的专用神经形态硬件来利用这些特性，使用从模拟电子学[7, 8]到标准CMOS ASIC [9-11]和完全可编程CPU [12, 13]的技术。
已经开发了许多用于直接训练SNN的技术。值得注意的是，反向传播可以应用于精确的脉冲时间[14-16]，或者每个脉冲神经元的不可微传递函数可以替换为"替代梯度"函数，从而允许通过时间进行反向传播[17, 18]或更多生物学合理的学习规则[19-21]应用于脉冲模型。然而，这些技术都还不能训练深度前馈SNN。对于这些模型，另一种方法是将经过充分训练的传统深度网络转换为SNN。最常见的方法是将深度网络中单元的激活转换为相同形状的SNN中脉冲神经元的发放率[22-24]。尽管使用这种转移方法已经获得了良好的推理性能，但在获得分类之前，需要对生成的SNN进行数百甚至数千个时间步骤的模拟。与每个ANN步骤相比，此开销抵消了每个SNN时间步骤的较低成本。另一种方法是转换深度网络模型，将ANN单元的激活编码在SNN模型[25, 26]中神经元的脉冲时间中。通过使用精确的脉冲时间，可以用更少的时间步骤表示激活，从而实现更具竞争力的推理速度。
脉冲神经网络可能最终部署在专门的神经形态硬件上，但研究人员仍然需要有效的工具来开发标准硬件上的转换和训练算法。有大量的SNN模拟器[27-30]，但其中大多数是为计算神经科学应用程序设计的，因此缺乏对卷积层和批处理等常见机器学习构建块的本地支持。因此，许多ML研究人员选择在PyTorch [35]、TensorFlow [36] 或JAX [37]等ML平台之上开发替代的基于脉冲的ML库[31-34]，以便他们可以继续利用这些平台提供的机器学习基础设施。这些SNN包各有优缺点，其中许多包专门模拟特定类别的模型。例如，SpykeTorch [34]是一个Python包，专门用于训练SNN模型，其中每个神经元每次刺激最多产生一个脉冲，而BindsNET [31]旨在成为一个更通用的SNN模拟包。
底层ML平台都通过对张量数据类型进行操作(例如矩阵乘法和归约)组装计算图来工作。这些操作被实现为高度优化的GPU核，为了确保它们完全占用GPU资源并抵消启动大量核的开销，操作在一个额外的张量维度上进行"批处理"。这是模拟前馈ANN的一种非常有效的方法，因为每个输入张量只需要一次遍历模型。然而，所采用的张量类型对稀疏性的支持很少，这意味着，为了在这些框架内模拟SNN，脉冲神经元的活动通常表示为具有许多零(无脉冲)和一些零(脉冲)的密集张量。然后使用基于标准矩阵乘法的方法通过网络的权重传播这些张量。因此，每个SNN时间步骤都需要与完整遍历等效ANN相同数量的算术运算。由于SNN模拟涉及处理多个时间步骤，因此当使用基于标准ML平台构建的库时，它们通常会比等效的ANN慢一个等于所需时间步数的因子。尽管如此，由于现代ML框架提供的高度优化的内核和强大的优化工具，这种方法已被证明在模拟小型模型方面很受欢迎。
在这项工作中，我们介绍了mlGeNN —— 一个将Keras中定义的前馈ANN模型转换为SNN模型的框架。它可在https://github.com/genn-team/ml_genn在线免费获得。在引擎盖下，mlGeNN不是使用基于张量的ML库来提供GPU加速，而是使用PyGeNN [38]——GeNN GPU加速SNN模拟器 [39]的Python接口，它之前在加速一系列计算神经科学模型[40, 41]方面表现出优异的性能。我们使用两种有效实现卷积连接的方法扩展了 GeNN。首先，我们使用以前用于模拟超大型计算神经科学模型[41]的"程序连接"方法直接实现2D卷积连接。其次，我们实施了一种针对二维卷积连接的稀疏结构量身定制的新并行策略。最后，我们为GeNN添加了批处理支持。通过这些改进，并且由于GeNN已针对连接的空间稀疏性和脉冲的时间稀疏性进行了优化，它可以显著提高SNN的模拟速度。然而，正如我们将在下面展示的那样，尽管与使用标准ML平台相比有了很大的改进，但SNN模拟所需的多个时间步骤意味着SNN模拟仍然比GPU硬件上的ANN慢。
2. Methods
2.1. GPU Architecture
2.2. Datasets
2.3 Models
2.4. Sparse connectivity for efficient convolutions
2.4.1. Procedural connectivity
2.4.2. Toeplitz connectivity
2.5. Batching
2.6. Rate-based conversion
2.7. Few-spike conversion
2.8. mlGeNN
3. Results
在本节中，我们探讨了使用第2.6节和第2.7节中描述的算法转换的SNN对准确性和推理延迟的影响。我们还将SNN的延迟与原始TensorFlow ANN以及对于VGG-16模型，与使用Lu和Sengupta[57]开发的修改以支持平均池化的BindsNet [31]模拟的SNN进行比较。
3.1. CIFAR-10
表1将第2.3节中描述的VGG-16和ResNet20模型的CIFAR-10数据集的准确性与Sengupta等人[24]以及St?ckl和Maass[26]先前的工作进行了比较。由于这两种模型都基于Sengupta等人[24]提供的描述，不出所料，两种人工神经网络模型都达到了相似的准确度水平。虽然mlGeNN确实支持Sengupta等人[24]使用的"spike-norm"基于发放率的转换算法，这里我们使用Diehl等人[22]的更简单的数据规范转换算法，经过修改以支持ResNet架构。这导致基于发放率的转换精度略低于Sengupta等人[24]报告的精度。但是，few-spike模型的表现要好得多。为了支持基于发放率的转换，我们在训练模型时不使用批量归一化和偏差。然而，St?ckl和Maass[26]不受这些限制的约束，因此能够训练出一个相当准确的ResNet-20模型，并将这种改进的性能转移到SNN模型中。
第2.6节和第2.7节中描述的两种转换算法都允许通过改变每个刺激呈现的模拟时间步数来"权衡"推理速度以获得准确性。当使用基于发放率的转换时，更多的时间步长(T)允许具有更细粒度的发放率代码，并且在使用few-spike转换时，额外的时间步长(K)会为二进制定点格式添加额外的"位"，从而减少量化误差。图4显示了VGG-16和ResNet-20模型在CIFAR-10数据集上训练并使用不同的 T 和 K 值转换为基于发放率和FewSpike SNN的准确性。当使用VGG-16模型时，分类精度在T < 2500或K < 10时显著下降(图4a)，而在使用ResNet-20模型时，分类性能仅在T < 1000或K < 8时显著下降(图4b)。ResNet-20模型的这种减少的"收敛时间"与Sengupta等人[24]提出的发现相匹配。
使用这些最佳 T 和 K 值，我们使用第2.4节中描述的不同连接算法在一系列批量大小上模拟了这两个模型。图5显示了使用VGG-16和ResNet-20模型对整个CIFAR-10测试集执行推理的最终时间。在批量大小为 1 时，第2.4.1节中描述的过程连接方法在所有转换算法和模型中表现最佳——ResNet-20 SNN实际上比TensorFlow ANN快3.5以上。然而，随着批量大小的增加，第2.4.2节中描述的Toeplitz矩阵算法表现最好，当使用ResNet-20模型时最快的FewSpike SNN仅比TensorFlow ANN慢2倍多一点，并且使用 VGG-16 模型时速度较慢大约慢12.5倍。程序连接算法的扩展性很差，因为在小批量大小时，每个神经元使用一个线程(而不是每个矩阵对角线)可以更好地利用GPU的并行性，随着批量大小的增加，内存延迟会由于用于更新突触后神经元输入的原子操作和严重合并的权重读取之间的高竞争而恶化。
因为使用基于发放率的算法转换的模型需要模拟比使用few-spike算法转换的模型更多的时间步骤，所以它们总体上更慢——即使是最快的模型也比TensorFlow ANN慢100多倍。然而，few-spike模型[26]使用的编码方案实际上比基于发放率的模型每时间步长大约多10个脉冲。因此，在基于发放率的模型中，每个时间步的脉冲活动要稀疏得多，这使得GeNN模拟单个时间步骤的速度明显快于ANN。请注意，每次评估的脉冲总数仍远低于基于发放率的转换的每次评估的总脉冲数，因为每次评估所需的时间步骤更少。
最后，将使用mlGeNN模拟的VGG-16模型与BindsNET [31]进行比较，我们可以看到，因为在底层，BindsNET使用与ANN相同的密集张量操作，如介绍中所讨论的，其性能与ANN乘以 T 相似，并且比基于最快发放率的GeNN VGG-16模型慢2.5倍以上。
3.2. ImageNet
表2将第2.3节中描述的VGG-16和ResNet-34模型的ImageNet数据集的准确性与Sengupta等人[24]提出的结果进行了比较。同样，这些模型基于Sengupta等人[24]提供的描述，所以人工神经网络模型达到了相似的准确度水平。根据我们在上一节中介绍的CIFAR-10结果，我们尚未在ImageNet上评估我们的基于发放率的转换算法，但我们的两个few-spike转换模型都优于Sengupta等人[24]报告的基于发放率的性能。
图6显示了使用VGG-16和ResNet-34模型对整个ImageNet测试集执行推理所需的时间。使用Toeplitz连接的模拟速度大约是两个模型中程序算法的两倍，但仍然比使用TensorFlow模拟的ANN慢100多倍。虽然起初可能有点令人惊讶，但这个结果可以在TensorFlow如何处理卷积的背景下理解。如第2.4节所述，TensorFlow在乘以Toeplitz矩阵或采用基于每个卷积层几何的Winograd或FFT方法之间切换。Lavin等人[50]使用这些方法导出了以下表达式，用于将Iwidth x Iheight x Ichan输入与Kheight x Kwidth x Ichan x Ochan核进行卷积所需的算术运算总数：
其中Nbatch是批量大小，m 是使用FFT或Winograd算法时输入被分成的tile大小的函数，并且与因子α, β, γ三种算法的不同部分中发生的操作有关。对于直接的基于矩阵乘法的方法，α = KwidthKheight并且β = γ = δ = 0，但对于Winograd和FFT方法，α, β, γ是值介于1和10之间的常数，具体取决于确切的配置。在使用大批量和大Iwidth、Iheight、Ichan和Ochan模拟的模型中，例如此处考虑的在ImageNet上运行的VGG和ResNet模型，β, γ, δ项相对较小，并且Winograd和FFT所需的操作数量基于方法本质上只是：
通过利用稀疏性，与直接矩阵乘法相比，GeNN减少了操作数(L)，乘以每个时间步骤脉冲的神经元的比例Pspike。此外，由于传播脉冲不需要乘法，因此应该有一些额外的优势。但是，如果基于Winograd和FFT的方法，α < PspikeKwidthKheight，则使用其中一种算法的ANN步骤将比GeNN当前脉冲传播算法的单个时间步骤更有效。因此，对于Winograd为2.25，TensorFlow在所考虑的模型中使用的FFT配置为1.81，并且Kwidth = Kheight = 3，SNN时间步骤比ANN步骤快需要Pspike < 0.2的稀疏性。对于此处报告的两种模型实现，Pspike ≈ 0.4，如果我们还考虑到我们为K = 10个时间步骤呈现每个刺激，TensorFlow在算术运算的数量上具有大约16个算法优势。此外，算术运算的数量只是故事的一半。cuDNN是一个高度优化的库，其中每个GPU架构的关键组件直接用SASS汇编语言编写，而GeNN基于 CUDA，还有很大的进一步优化空间，特别是消除对全局内存中原子操作的需求，以及进一步优化指标计算。
4. Discussion
在本文中，我们介绍了mlGeNN，这是一个使用GeNN GPU加速SNN模拟器的脉冲神经网络机器学习库。 我们已经描述了mlGeNN如何使用标准算法[22, 24, 26]将经过训练的ANN架构转换为不同的SNN实现，以及推理的速度和准确性如何与原始ANN模型和其他库模拟的SNN模型进行比较。我们发现准确性与原始出版物中有关ANN到SNN转换的报告中描述的一样，而推理速度呈现出非常复杂的情况。虽然在大多数考虑的示例中，mlGeNN比BindsNET [31]快，并且在某些示例中与ANN速度相当，但对于其他示例，它要慢几个数量级。
如第3.2节所述，在相同硬件上运行的ANN和SNN模型之间的性能差距来自SNN模型必须运行每个刺激的时间步骤、专用ANN卷积算法的算法优势以及优化程度和与GeNN [39]等SNN库相比，cuDNN [51]的成熟度。
在这项工作中——除了 2.6 和 2.7 节中描述的限制——我们只考虑了使用标准32位浮点权重和激活的模型转换。 然而，在尽可能少的SNN时间步骤中表示ANN激活——使用发放率代码或由few-shot方法使用的定点脉冲代码——与量化ANN激活以使用低精度数据类型进行推理的问题类似。因此，将量化感知训练[58]或微调[5]中的一些技术应用于我们的ANN激活可以使每个刺激的SNN时间步数最小化，从而减少ANN和SNN性能之间的差距。
虽然FewSpike表示每个刺激需要更少的SNN时间步骤，但在我们在第3.2节评估的ImageNet规模模型中，它们导致非常高的发放率，这使得每个时间步骤都比ANN更新慢。一种解决方案是通过使用L1正则化训练ANN来"稀疏"神经元激活，因此更多的激活为零，这由没有脉冲表示。或者，可以使用用更少位表示更常见激活值的二值代码，尽管这会给模型的神经元或突触增加额外的解码复杂性。
我们在第2.4.2节中描述的新的"Toeplitz"连接方案明显快于Knight和Nowotny[40]描述的通用稀疏连接以及Knight等人[38]描述的通用"程序"连接。但是，仍然需要显著程度的稀疏性才能匹配ANN更新所需的操作数量。这一特性不仅会影响像GeNN这样的基于软件的SNN模拟器，而且对于旨在与非脉冲ML加速器竞争的基于尖峰的神经形态芯片的设计者来说，也是一个重要的考虑因素。然而，除了我们之前讨论过的更新ANN所需操作数量的理论差异之外，还有很大的空间可以进一步优化我们的实现。我们当前的实现使用"即发即弃"原子添加操作将每一层的输出累积到全局内存中。尽管这些操作完全在现代GPU的L2缓存中处理，因此延迟相对较低，但通过使用Bautembach等人[53]描述的方法将更新累积到共享内存中，这可以进一步减少。此外，使用NVIDIA Nsight计算分析器进行分析表明，虽然该算法目前主要受计算限制，但大部分算术运算涉及索引计算，而不是实际处理卷积。这是因为我们的单程序多数据代码生成器(在Knight等人[38]中描述)将重用相同的生成代码来模拟所有具有Toeplitz连通性的层，这意味着CUDA编译器将无法优化除法运算2.4.2节中介绍的算法所需的指数计算。当前的 GPU 架构没有对整数除法的硬件支持，因此这会导致大量多余的操作来在软件中实现这些操作，但由于除数在构建时都是已知的，因此可以用单个整数乘加操作代替[59]。
总之，mlGeNN框架为GeNN生态系统添加了一个新元素，允许将GeNN用于机器学习工作流程，将经过充分训练的深度ANN转换为SNN以进行推理。这是建立在GeNN之上的完全基于SNN的机器生态系统的第一步，未来也将允许直接训练SNN。最终，我们希望这将成为开发基于SNN的算法以部署在标准和神经形态硬件上的理想工具。