研究人员表示,他们将边缘训练看成是一个优化问题,从而发现了最优调度,以实现下的能耗最小给定的内存预算。目前,深度学习模型已经广泛部署在智能手机、嵌入式平台等边缘设备上进行推理。其中,训练仍主要在大型云服务器上进行,配备GPU等高通量加速器。集中式云训练模型需要将照片和击键等敏感数据从边缘设备传输到云端,这会牺牲用户隐私并产生额外的数据移动成本。图注:推特@ShishirPatil因此,为了让用户在不牺牲隐私的情况下个性化他们的模型,联邦学习等设备上的训练方法不需要将数据集成到云端,也可以进行本地训练更新。这些方法已部署在Google的Gboard键盘上以个性化键盘建议,iPhone也使用这些方法来改进自动语音识别。同时,当前的设备端训练方法不支持训练现代架构和大型模型。在边缘设备上训练更大的模型是不可行的,主要是因为有限的设备内存无法存储反向传播激活。ResNet-50的单次训练迭代需要的内存是推理的200多倍。先前工作提出的策略包括分页到辅助内存和重新实现以减少云训练的内存占用。然而,这些方法会显着增加整体能耗。与分页方法相关的数据传输通常比重新计算数据需要更多的能量。随着内存预算的缩减,重新实现会以O(n^2)的速度增加能耗。在加州大学伯克利分校最近的一篇论文中,几位研究人员表明分页和重新实现是高度互补的。通过重新实现简单操作,同时将复杂操作的结果分页到闪存或SD卡等辅助存储,他们能够以最小的能耗扩展有效内存容量。而且,通过结合这两种方法,研究人员还证明了可以在移动级边缘设备上训练BERT等模型。通过将边缘训练视为优化问题,他们发现了在给定内存预算下实现最低能耗的最佳计划。论文地址:https://arxiv.org/pdf/2207.07697.pdf项目主页:https://poet.cs.berkeley.edu/GitHub地址:https://github.com/shishirpatil/poet研究者提出了POET(PrivateOptimalEnergyTraining),这是一种在内存受限的边缘设备上进行现代神经网络能量优化训练的算法。其架构如下图1所示。鉴于为反向传播缓存所有激活张量的成本过高,POET优化了分页和激活的重新实现,将内存消耗减少了两倍。他们将边际训练问题重新表述为整数线性规划(ILP),并发现求解器可以在10分钟内将其求解至最优。图注:POET优化了SOTA机器学习模型在边缘设备上的训练。对于部署在真实世界边缘设备上的模型,训练会在边缘设备空闲且计算周期可用时进行,例??如,GoogleGboard在手机充电时安排模型更新。因此,POET也包含了严格的训练约束。给定内存限制和训练时期的数量,POET生成的解决方案也满足给定的训练期限。此外,研究人员利用POET开发了一个综合成本模型,并证明它对于现有的开箱即用架构在数学上是合理的(即不做近似)。该论文的第一作者ShishirPatil在一段演示视频中表示,POET算法可以训练任何在智能手机等商业边缘设备上需要大量内存的SOTA模型。他们还成为第一个展示在智能手机和ARMCortex-M设备上训练SOTA机器学习模型(例如BERT和ResNet)的研究团队。集成分页和重新实现重新实现和分页是减少大型SOTAML模型内存消耗的两种技术。在重新实现中,激活张量一旦不再需要就会被删除,最常见的是在前向传播过程中。这释放了宝贵的内存,可用于存储后续层的激活。当再次需要删除的张量时,该方法会根据沿袭规定的其他相关激活重新计算。分页,也称为卸载,是一种减少内存的补充技术。在分页中,将不需要立即使用的激活张量从主内存调出到辅助存储,例如闪存或SD卡。当再次需要张量时,它会被调出。图2显示了八层神经网络的执行时间线。沿x轴,每个单元对应神经网络的每一层(共8层L8)。Y轴表示一个纪元内的逻辑时间步长。图中被占用的单元格(用颜色填充)表示在相应时间步执行的操作(前向/后向传播计算、重新实现或分页)。例如,我们可以看到L1的激活是在第一个时间步(T1)计算的。在T2和T3,分别计算L2和L3的激活。假设层L2和L3恰好是内存密集型但计算成本低廉的操作,例如非线性(tanH、ReLU等),那么重新实现成为最佳选择。我们可以删除激活({T3,L2},{T4,L3})以释放内存,并在向后传递({T14,L3},{T16,L2})期间需要时重新实现它们。假设L5和L6层是计算密集型操作,如卷积、密集矩阵乘法等。对于此类操作,重新实现会导致运行时间和能量增加,这种方法不是最优的。对于这些层,最好将激活张量分页到辅助存储({T6,L5},{T7,L6}),并在需要时分页到({T10,L6},{T11,L5})。分页的一个主要优点是,根据内存总线占用情况,它可以流水线化以隐藏延迟。这是因为现代系统具有DMA(直接内存访问)功能,它将激活张量从辅助存储移动到主内存,同时计算引擎并行运行。例如,在时间步T7,可以同时调用L6和计算L7。然而,重新实现是计算密集型的并且不能并行化,这导致运行时间增加。例如,我们必须使用时间步T14重新计算L3,延迟反向传播的其余部分执行。POET这项研究介绍了POET,这是一种用于深度神经网络的图级编译器,它重写了大型模型的训练DAG,以适应边缘设备的内存限制,同时保持能源效率。POET是硬件感知的,它首先跟踪前向和后向传递的执行以及相关的内存分配请求、运行时间以及每个操作的内存和能量消耗。这种针对给定硬件的每个工作负载的细粒度分析只发生一次,是自动化的、廉价的,并为POET提供了最准确的成本模型。POET然后生成一个可以有效求解的混合整数线性规划(MILP)。POET优化器搜索有效的重新实现和分页,以最大限度地减少内存绑定的端到端能耗。然后利用得到的schedule生成新的DAG,在边缘设备上执行。虽然MILP是在商品硬件上解决的,但发送到边缘设备的调度只有几百个字节,因此内存效率非常高。对于计算成本低但内存密集的操作,重新实现最有效。但是,分页最适合计算密集型操作,在这些操作中重新实现会产生大量能量开销。POET共同考虑了集成搜索空间中的重新实现和分页。本文的方法可以推广到复杂和现实的架构,POET优化器算法如下。本研究在优化问题中引入了一个新的目标函数来最小化计算、page-in和page-out的综合能耗,新的分页和重新实现能耗的目标函数为:其中Φ_compute、Φ_pagein和Φ_pageout分别代表每个节点在计算、页入和页出时所消耗的能量。POET被重新实现以根据图形的哪些节点(k)以及在每个时间步(t)哪些节点是页入或页出来输出DAG计划。实验结果在评估POET时,研究人员试图回答三个关键问题。首先,POET在不同模型和平台上降低了多少能源消耗?其次,POET如何从混合分页和重新实现策略中获益?最后,POET如何适应不同的运行时预算?研究人员在下表2中列出了四种不同的硬件设备,分别是ARMCortexM0MKR1000、ARMCortexM4Fnrf52840、A72RaspberryPi4B+和NvidiaJetsonTX2。POET是完全硬件感知的,并且依赖于细粒度的分析。下图3显示了单个训练epoch的能量消耗,每一列对应不同的硬件平台。研究人员发现,POET在所有平台上生成节能计划(Y轴),同时减少峰值内存消耗(X轴)并满足时间预算。在下面的图5中,我们在A72上训练ResNet-18时对POET和Capuchin进行了基准测试。在RAM预算减少的情况下,Capuchin消耗的能量比全内存基线多73%到141%。相比之下,POET产生的能量不到所消耗能量的1%。这种趋势适用于所有测试的架构和平台。在表3中,该研究在Nvidia的JetsonTX2上训练ResNet-18时对POET和POFO进行了基准测试。研究发现,POET发现了一种集成的重新实现和分页,可以将峰值内存消耗降低8.3%,并将吞吐量提高13%。这展示了POET的MILP求解器的实力,它能够在更大的搜索空间上进行优化。虽然POFO仅支持线性模型,但POET可以推广到非线性模型,如图3所示。图4突出显示了POET在不同时间约束下采用集成策略的好处。对于每个运行时间,下图绘制了总能耗。
