基于神经网络中静态计算图的内存优化方法及电子设备

技术领域

本发明涉及新一代信息技术领域，特别是涉及神经网络技术领域。

背景技术

随着以神经网络为主要内容的深度学习算法的不断进步，人工智能算法成功地广泛应用到了现实生活中的许多重要领域，包括图像处理、语言处理、文字处理等。作为深度神经网络的重要组成部分，许多不同的激活函数被用于不同类型的任务中，从而不断提高模型精度。比如说，对大多数视觉任务而言，使用Mish、Swish或GELU替代ReLU作为激活函数能够在ResNet架构的神经网络中取得1％的精度提升。另一方面，在自然语言处理领域中，GELU激活函数通常优于其他候选的激活函数。

然而，新型的激活函数通常有着复杂的结构与高昂的内存占用，尤其是在结合了未充足优化的实现之后。有些情况下，激活函数的内存占用在梯度计算过程中会超过卷积层与全连接层在前向过程中保存的中间变量成为最大的内存消耗方。这极大阻碍了新型激活函数在GPU等显存非常有限的主流神经网络训练加速硬件上的应用。

新型激活函数的高内存占用主要是因为主流的神经网络框架采用后向传播算法作为训练过程中自动微分的默认实现。后向自动微分需要保存前向过程中激活函数的中间变量，带来了额外的内存占用。而现有的适用于后向自动微分的内存优化方法通常会引入额外的计算开销，从而降低了神经网络训练速度。另一方面，采用前向传播算法实现自动微分能够降低新型激活函数的内存占用，但在卷积层、全连接层等神经网络主流结构中前向自动微分的计算效率又远低于后向自动微分，因此前向自动微分无法完全替代后向自动微分。

重计算(recomputation)是一种主流的基于后向自动微分的内存优化技术。后向自动微分将神经网络训练过程分为前向传播和后向传播两个阶段，在后向传播过程中，后向自动微分需要访问前向传播过程中产生的中间变量，因此在后向传播阶段对应层的梯度没有计算完之前，对应的中间变量必须一直保留，这些中间变量构成了后向自动微分过程中的主要内存占用。重计算技术为了减少不必要的中间变量占用，主动丢弃一部分经过简单计算就能恢复出来的中间变量，在后向传播阶段再根据需要重新计算对应的中间变量，从而节约了内存占用。在这个过程中，因为激活函数需要的计算资源远低于卷积层、全连接层等神经网络结构，绝大多数激活函数生成的中间变量会被丢弃。

这种方法存在的主要问题是引入了额外的计算开销。为了在后向传播阶段恢复出激活函数之后的中间变量，激活函数会被反复调用，如果新型激活函数中含有一些开销较大的计算步骤或者相对复杂的计算流程，那么对训练速度的负面影响便不可小觑。

采用前向自动微分作为神经网络框架默认的自动微分实现是一种有效的内存优化技术。与后向自动微分不同，前向自动微分没有后向传播过程，训练神经网络需要的梯度在前向传播过程中计算。由于不需要将梯度计算拆分到后向传播阶段，前向自动微分不需要保留前向传播阶段的中间变量，从而有效减少了训练过程中的内存占用。

这种方法在进行快速的量化的同时无法保证神经网络识别精度。这种方法的主要问题是前向自动微分采用的求导函数比后向自动微分的求导函数在卷积层、全连接层等占据主要计算资源的神经网络结构中计算开销更高，从而导致默认采用前向自动微分的训练过程耗时更长，无法被主流神经网络框架采用。

由此可见，当前已有的神经网络训练内存优化方法在处理激活函数上存在诸多问题，无法在减少内存占用的同时不引入额外的计算开销。这给超大模型训练的应用造成了严重阻碍，尤其是在大模型快速普及的背景下。因此如何减低神经网络训练过程中的内存占用，提高神经网络的训练速度，降低神经网络训练的资源消耗成为本领域技术人员亟待解决的技术问题。

发明内容

鉴于以上所述现有技术的缺点，本发明的目的在于提供一种基于神经网络中静态计算图的内存优化方法及电子设备，用于减低神经网络训练过程中的内存占用，提高神经网络的训练速度，降低神经网络训练的资源消耗。

为实现上述目的及其他相关目的，本发明提供一种基于神经网络中静态计算图的内存优化方法，包括：识别神经网络静态计算图中前向兼容算子；消除与所述前向兼容算子对应的与后向自动微分相关的子图；生成与所述前向兼容算子对应的前向求导子图，并输出当前的神经网络静态计算图。

于本发明的一实施例中，所述识别神经网络静态计算图中前向兼容算子的方式包括以下任一种或多种组合：在当前算子为一元算子时，确定所述当前算子为前向兼容算子；将具有一个非常量输入的n元算子确定为前向兼容算子；将所有输入来自同一上游算子的n元算子确定为前向兼容算子。

于本发明的一实施例中，将由前向兼容算子组成的组合算子确定为前向兼容算子。

于本发明的一实施例中，还包括：将由互相连通的前向兼容算子组成的子图标记为前向兼容子图。

于本发明的一实施例中，所述消除与所述前向兼容算子对应的与后向自动微分相关的子图包括：消除与前向兼容子图对应的后向子图间的连接；消除与前向兼容子图对应的后向子图。

于本发明的一实施例中，还包括：消除与前向兼容子图对应的后向阶段的求导算子。

于本发明的一实施例中，所述生成与所述前向兼容算子对应的前向求导子图中，基于前向自动微分的求导方法生成对应的前向求导子图，并将生成的所述前向求导子图插入所述静态计算图中。

于本发明的一实施例中，还包括：对于每个前向兼容子图，将其对应的前向求导子图的输出与后向图中对应的梯度结合生成在后向图中的求导算子，从而修复后向图。

为实现上述目的及其他相关目的，本发明还提供一种计算机可读存储介质，其上存储有计算机程序，该程序被处理器执行时实现如上所述的基于神经网络中静态计算图的内存优化方法。

为实现上述目的及其他相关目的，本发明还提供一种电子设备，包括存储器，用于存储计算机程序；处理器，用于运行所述计算机程序以实现如上所述的基于神经网络中静态计算图的内存优化方法。

如上所述，本发明的基于神经网络中静态计算图的内存优化方法及电子设备具有以下有益效果：

1、本发明能够在不引入额外计算开销的情况下，通过自动优化静态计算图有效降低大部分神经网络激活函数的内存占用。

2、本发明为大模型在内存受限设备上进行训练提供了有效支持，为新型激活函数的大规模应用铺平了道路，为大型神经网络基于静态计算图的训练节约了大量内存资源。

附图说明

为了更清楚地说明本发明实施例中的技术方案，下面将对实施例描述中所需要使用的附图作简单地介绍，显而易见地，下面描述中的附图仅仅是本发明的一些实施例，对于本领域普通技术人员来讲，在不付出创造性劳动的前提下，还可以根据这些附图获得其他的附图。

图1显示为本申请一实施例中的基于神经网络中静态计算图的内存优化方法的原理示意图；

图2显示为本申请一实施例中的基于神经网络中静态计算图的内存优化方法中前向兼容算子与非前向兼容算子示意图；

图3显示为本申请一实施例中的基于神经网络中静态计算图的内存优化方法中静态计算图优化算法示意图；

图4显示为本申请一实施例中的基于神经网络中静态计算图的内存优化方法中整体实施过程示意图；

图5显示为本申请一实施例中的基于神经网络中静态计算图的内存优化方法的逻辑流程图；

图6显示为本申请一实施例中的电子设备的原理框图。

具体实施方式

以下通过特定的具体实例说明本发明的实施方式，本领域技术人员可由本说明书所揭露的内容轻易地了解本发明的其他优点与功效。本发明还可以通过另外不同的具体实施方式加以实施或应用，本说明书中的各项细节也可以基于不同观点与应用，在没有背离本发明的精神下进行各种修饰或改变。需说明的是，在不冲突的情况下，以下实施例及实施例中的特征可以相互组合。

本实施例提出了一种在后向自动微分中嵌套前向自动微分的内存优化方法，该方法能够自动识别出神经网络中需要替换为前向自动微分的激活函数，并在主流神经网络框架的静态计算图中自动替换对应的子图，从而在不引入额外计算开销的同时，降低神经网络中激活函数的内存占用，自动节省神经网络训练过程中的资源消耗，允许在同样的硬件条件下训练更大更复杂的神经网络。

具体地，本实施例的目的在于提供一种基于神经网络中静态计算图的内存优化方法及电子设备，用于减低神经网络训练过程中的内存占用，提高神经网络的训练速度，降低神经网络训练的资源消耗。

以下将详细阐述本发明的基于神经网络中静态计算图的内存优化方法及电子设备的原理及实施方式，使本领域技术人员不需要创造性劳动即可理解本发明的基于神经网络中静态计算图的内存优化方法及电子设备。

实施例1

如图1所示，本实施例提供一种基于神经网络中静态计算图的内存优化方法，包括：

步骤S100，识别神经网络静态计算图中前向兼容算子；

步骤S200，消除与所述前向兼容算子对应的与后向自动微分相关的子图；

步骤S300，生成与所述前向兼容算子对应的前向求导子图，并输出当前的神经网络静态计算图。

以下对本实施例的基于神经网络中静态计算图的内存优化方法进行详细说明。

本实施例的基本原理在于神经网络模型中，前向自动微分相对于后向自动微分在大部分激活函数上计算效率更高、内存占用更少。在阐述该基本原理之前，需要对前向自动微分与后向自动微分的求导过程进行描述。

对于一个函数f：

其中，M表示输出向量的长度，N表示输入向量的长度，J

前向自动微分需要计算雅可比-向量积(Jacobian-vector product，JVP)，表示为J

以Swish激活函数为例，Swish激活函数是一个由Mul和Sigmoid两个操作组合成的激活函数，其计算过程可表示为：

v

其中v

在Swish激活函数上，前向自动微分需要计算v

v′

同理可得，对后向自动微分而言，需要计算对应的向量-雅可比积：

其中

从上述对前向自动微分和后向自动微分求导的推导结果可以看出，在f：

接下来本实施例描述前向自动微分嵌入后向自动微分过程的基本原理。假设神经网络计算图

其中，

其中g为上游梯度。因此本实施例可以使用前向自动微分在前向传播阶段计算由v

基于上述思想，以下将详细介绍本实施例的实施细节。

步骤S100，识别神经网络静态计算图中前向兼容算子。

为了识别出静态计算图中有哪些算子可以使用嵌套前向自动微分，如图2所示，本实施例将计算图中的算子分为两类：前向兼容算子与非前向兼容算子。

于本实施例中，所述识别神经网络静态计算图中前向兼容算子的方式包括以下任一种或多种组合：

1)在当前算子为一元算子时，确定所述当前算子为前向兼容算子。

所有一元算子(unary operation)都是前向兼容算子。一元算子可表示为f：

2)将具有一个非常量输入的n元算子确定为前向兼容算子。

所有只有一个非常量输入的n元算子(n-ary operation)都是前向兼容算子。常见的n元算子为二元算子(binary operation)，例如图2中的节点2和节点3都是二元算子.由于n个输入中只有一个不是常量，因此只有一个非常量输入的n元算子可以视作一种特殊的一元算子。

3)将所有输入来自同一上游算子的n元算子确定为前向兼容算子。

所有输入来自同一上游算子的n元算子都是前向兼容算子。例如图2中的节点2虽然有两个输入，但两个输入都最终来自同一上游节点，因此它是一个前向兼容节点。

所有不满足上述条件的基础算子都是非前向兼容算子。

对于由有限个基础算子组成的组合算子，使用以下算法流程自动识别前向兼容算子：

1)所有只由前向兼容算子组成的组合算子都是前向兼容算子。

2)所有不满足上述条件的组合算子都是非前向兼容算子。

即于本实施例中，将由前向兼容算子组成的组合算子确定为前向兼容算子。

于本实施例中，还包括：将由互相连通的前向兼容算子组成的子图标记为前向兼容子图。

步骤S200，消除与所述前向兼容算子对应的与后向自动微分相关的子图。

于本实施例中，所述消除与所述前向兼容算子对应的与后向自动微分相关的子图包括：

1)消除与前向兼容子图对应的后向子图间的连接；

2)消除与前向兼容子图对应的后向子图。

于本实施例中，还包括：消除与前向兼容子图对应的后向阶段的求导算子。

步骤S300，生成与所述前向兼容算子对应的前向求导子图，并输出当前的神经网络静态计算图。

于本实施例中，所述生成与所述前向兼容算子对应的前向求导子图中，基于前向自动微分的求导方法生成对应的前向求导子图，并将生成的所述前向求导子图插入所述静态计算图中。

于本实施例中，还包括：对于每个前向兼容子图，将其对应的前向求导子图的输出与后向图中对应的梯度结合生成在后向图中的求导算子，从而修复后向图。

如图3所示，本实施例中，静态计算图优化算法的实施过程如下：

1)使用前向兼容算子自动识别算法标记出前向阶段中所有前向兼容算子，将由互相连通的前向兼容算子组成的子图标记为前向兼容子图。

2)对于每个前向兼容子图，消除其对应的后向阶段的求导算子，消除其前向子图与后向子图间用于保存中间变量的连接。

3)对于每个前向兼容子图，基于前向自动微分的求导方法生成对应的前向求导子图。

4)对于每个前向兼容子图，将其对应的前向求导子图的输出与后向图中对应的梯度结合生成在后向图中的求导算子，从而修复后向图。

如图4所示，本实施例的基于神经网络中静态计算图的内存优化方法分为前向兼容算子自动识别、原后向子图消除、新前向求导子图生成三部分。首先使用前向兼容算子自动识别算法标记出所有的前向兼容子图作为后续优化的操作范围。对于每一个前向兼容子图，需要将其从使用后向自动微分切换为使用前向自动微分。切换过程的第一阶段是消除原有的与后向自动微分有关的子图，这包括消除前向兼容子图与对应的后向子图间的连接与消除整个后向子图两个步骤。之后，通过生成前向兼容子图对应的前向求导子图，成功将前向兼容子图切换为使用前向自动微分。经过上述优化后，最终生成了内存优化后的静态计算图作为算法输出。

图5显示为本实施例中的基于神经网络中静态计算图的内存优化方法的逻辑流程图；如图6所示，本实施例中的基于神经网络中静态计算图的内存优化方法的逻辑流程如下：

1)选择待优化的前向算子。

2)使用前向兼容算子自动识别算法判断该算子是否为前向兼容算子。

3)如果为前向兼容算子，则进行后续优化；如果不是前向兼容算子，选择下一个前向算子进行优化。

4)消除前向兼容算子和对应的后向子图间的连接。

5)消除前向兼容算子对应的后向子图。

6)生成前向兼容算子对应的前向求导算子，并插入静态计算图中。

7)在遍历了所有前向算子之后，优化过程已完成，输出优化后的静态计算图并结束。

本实施例中，用户与软件层的交互：用户指定系统所需要优化的神经网络静态计算图。

本实施例中，运行逻辑包括：

1)根据用户提供的神经网络静态计算图识别出需要优化的子图；

2)将待优化子图从后向自动微分切换为前向自动微分，从而生成优化后的神经网络静态计算图。

本实施例是一种在后向自动微分中嵌套前向自动微分的内存优化方案，能够在不引入额外计算开销的情况下，通过自动优化静态计算图有效降低大部分激活函数的内存占用。本实施例为大模型在内存受限设备上进行训练提供了有效支持，为新型激活函数的大规模应用铺平了道路，为大型神经网络基于静态计算图的训练节约了大量内存资源。

实施例2

如图6所示，本实施例提供一种电子设备101，所述电子设备101包括：处理器1001及存储器1002；所述存储器1002用于存储计算机程序；所述处理器1001用于执行所述存储器1002存储的计算机程序，以使所述电子设备101执行如实施例1中基于神经网络中静态计算图的内存优化方法的步骤。由于基于神经网络中静态计算图的内存优化方法的步骤的具体实施过程已经在实施例1中进行了详细说明，在此不再赘述。

处理器1001为(Central Processing Unit，中央处理器)。存储器1002通过系统总线与处理器1001连接并完成相互间的通信，存储器1002用于存储计算机程序，处理器1001用于运行计算机程序，以使所述处理器1001执行所述的基于神经网络中静态计算图的内存优化方法。存储器1002可能包含随机存取存储器(Random Access Memory，简称RAM)，也可能还包括非易失性存储器(non-volatile memory)，例如至少一个磁盘存储器。

此外，本实施例还提供一种计算机可读存储介质，其上存储有计算机程序，所述计算机程序被处理器1001执行时实现实施例1中所述的基于神经网络中静态计算图的内存优化方法。实施例1已经对所述基于神经网络中静态计算图的内存优化方法进行了详细说明，在此不再赘述。

本领域普通技术人员可以理解：实现上述各方法实施例的全部或部分结构可以通过计算机程序相关的硬件来完成。前述的计算机程序可以存储于一计算机可读存储介质中。该程序在执行时，生成上述基于神经网络中静态计算图的内存优化方法；而前述的存储介质包括：ROM、RAM、磁碟或者光盘等各种可以存储程序代码的介质。

综上所述，本发明能够在不引入额外计算开销的情况下，通过自动优化静态计算图有效降低大部分神经网络激活函数的内存占用；本发明为大模型在内存受限设备上进行训练提供了有效支持，为新型激活函数的大规模应用铺平了道路，为大型神经网络基于静态计算图的训练节约了大量内存资源。所以，本发明有效克服了现有技术中的种种缺点而具高度产业利用价值。

上述实施例仅例示性说明本发明的原理及其功效，而非用于限制本发明。任何熟悉此技术的人士皆可在不违背本发明的精神及范畴下，对上述实施例进行修饰或改变。因此，举凡所属技术领域中包括通常知识者在未脱离本发明所揭示的精神与技术思想下所完成的一切等效修饰或改变，仍应由本发明的权利要求所涵盖。