CN117172289A - 张量切分方法、装置以及电子设备 - Google Patents

Abstract

本申请公开一种张量切分方法及装置，张量切分方法包括：基于神经网络的输入数据，获取神经网络对该输入数据进行计算时的计算图以及张量数据的张量信息；根据计算图的拓扑顺序对神经网络的张量数据进行内存分配，并从计算图中获取当前节点和当前张量数据，当前张量数据的尺寸超过当前内存的可分配空间；从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组，当前张量切分组包括至少2个张量数据；对当前张量切分组内的张量数据进行切分以使当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间。本申请可以对大的张量数据进行切分以满足内存空间要求，减少对外部存储器的访问，提高计算效率。

Description

张量切分方法、装置以及电子设备

技术领域

本发明涉及人工智能技术领域，更具体地，涉及一种张量切分方法、装置以及电子设备。

背景技术

神经网络模型在运行过程中按照计算图的拓扑顺序依次执行各个算子的计算。相应地，在编译过程中按照拓扑顺序对各个算子的张量数据进行内存分配。随着人工智能(Artificial Intelligence，AI)技术的发展，神经网络模型也越来越复杂，神经网络模型的张量数据在内存分配过程中可能超出内存的最大容量导致内存分配失败而无法运行。

相关技术中，通过将算子的输出张量数据存储在外部存储器，然后将内存空间内存储的算子的输入张量数据以及产生的输出数据先删除以释放内存空间。最后，再从外部存储器中获取该算子的输出数据，并作为下一算子的输入数据。但是外部存储器的读取速度远远低于内存的读取速度，会影响神经网络模型的计算效率。

因此，在神经网络模型编译过程中，如何对算子进行张量切分以使各个算子的中间数据存储在内存中，减少外部存储器的访问，是当前亟待解决的技术问题。

发明内容

鉴于上述问题，本发明的目的在于提供一种张量切分方法、装置以及电子设备，可以使神经网络的张量数据的读取和存放均在内存中，减少对外部存储器的访问，提高计算效率。

根据本发明的第一方面，提供一种神经网络的张量切分方法，包括：基于神经网络的输入数据，获取神经网络对该输入数据进行计算时的计算图以及张量数据的张量信息，其中，所述计算图包括多个计算节点以及各个计算节点之间的张量数据；根据计算图的拓扑顺序对神经网络的张量数据进行内存分配，并从所述计算图中获取当前节点和当前张量数据，其中，所述当前张量数据的尺寸超过当前内存的可分配空间；从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组，其中，所述当前张量切分组包括至少2个张量数据；对当前张量切分组内的张量数据进行切分以使当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间。

优选地，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组包括：当神经网络具有分支结构时，从当前节点向前反向搜索至起始节点，所述起始节点为与当前节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组，所述终止节点为起始节点之后分支的根节点，所述根节点为所有分支的合并节点，当前节点分组包括起始节点、终止节点以及起始节点和终止节点之间的中间节点；将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

优选地，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组还包括：判断当前节点是否为起始节点或者终止节点；若当前节点为起始节点或者终止节点，从当前节点分组的起始节点继续反向向前搜索以更新当前节点分组的起始节点，更新后的起始节点为当前起始节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组；当当前节点不为起始节点或终止节点时，将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组；重复以上步骤直至当前节点不为起始节点或终止节点。

优选地，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组还包括：遍历当前节点分组内是否存在不支持切分的中间节点；所述当前节点分组内存在不支持切分的中间节点，从不支持切分的中间节点处将当前节点分组进行截断以分成至少2个节点分组；当所述节点分组均为支持切分的中间节点，将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

优选地，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组还包括：当所述当前节点分组的中间节点的个数超过预设阈值时，将所述当前节点分组内去除起始节点与终止节点之间距离近的分支结构以获取缩小后的所述当前节点分组。

优选地，对当前张量切分组内的张量数据进行切分包括：遍历所有切分策略对当前张量切分组的张量数据进行切分；判断当前张量切分组内的切分后的张量数据的尺寸是否超过内存的可分配空间；当当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间时进行内存分配。

根据本发明的第二方面，提供一种张量切分装置，包括：获取单元，用于基于神经网络的输入数据，获取神经网络对该输入数据进行计算时的计算图以及张量数据的张量信息，其中，所述计算图包括多个计算节点以及各个计算节点之间的张量数据；分配单元，用于根据计算图的拓扑顺序对神经网络的张量数据进行内存分配，并从所述计算图中获取当前节点和当前张量数据，其中，所述当前张量数据的尺寸超过当前内存的可分配空间；搜索单元，用于从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组，其中，所述当前张量切分组包括至少2个张量数据；切分单元，用于对当前张量切分组内的张量数据进行切分以使当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间；其中，所述搜索单元还用于当神经网络具有分支结构时，从当前节点向前反向搜索至起始节点，所述起始节点为与当前节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组，所述终止节点为起始节点之后分支的根节点，所述根节点为所有分支的合并节点，当前节点分组包括起始节点、终止节点以及起始节点和终止节点之间的中间节点；将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

优选地，所述搜索单元还用于判断当前节点是否为起始节点或者终止节点；若当前节点为起始节点或者终止节点，从当前节点分组的起始节点继续反向向前搜索以更新当前节点分组的起始节点，更新后的起始节点为当前起始节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组；当当前节点不为起始节点或终止节点时，将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组；重复以上步骤直至当前节点不为起始节点或终止节点。

根据本发明的第三方面，提供一种电子设备，包括存储器、处理器以及存储在所述存储器中并可在所述处理器上运行的计算机程序，所述处理器执行所述计算机程序时实现上述所述的张量切分方法。

根据本发明的第四方面，提供一种计算机可读存储介质，所述计算机可读存储介质存储有计算机程序，所述计算机程序被处理器执行时实现上述所述的张量切分方法。

本发明提供的张量切分方法、装置以及电子设备，通过反向搜索获取前节点分组以及当前张量切分组，并对当前张量切分组内的张量数据进行切分以使神经网络的张量数据的读取和存放均在内存中，可以对大的张量数据进行切分以满足内存空间要求，减少对外部存储器的访问，提高计算效率。

进一步地，当所述当前节点分组存在不支持切分的中间节点时，从不支持切分的中间节点处将所述当前节点分组划分成至少2个节点分组，解决节点分组内有些节点的输入输出张量无法切分的问题。

进一步地，当所述节点分组的中间节点超过预设阈值时，对所述节点分组进行缩小处理以获取缩小后的所述节点分组，减少数据重叠和重复计算，进一步提高计算效率。

附图说明

通过以下参照附图对本发明实施例的描述，本发明的上述以及其他目的、特征和优点将更为清楚，在附图中：

图1示出本申请的一些实施例的一种计算图10的示意图；

图2示出本申请的一些实施例的另一种计算图20的示意图；

图3示出本申请的一些实施例的另一种计算图30的示意图；

图4示出本申请实施例提供的神经网络的张量切分方法的流程图；

图5示出本申请实施例提供的步骤S130的流程图；

图6示出本申请实施例提供的神经网络的张量切分装置的结构示意图；

图7示出本申请实施例提供的电子设备的结构示意图。

具体实施方式

以下将参照附图更详细地描述本发明的各种实施例。在各个附图中，相同的元件采用相同或类似的附图标记来表示。为了清楚起见，附图中的各个部分没有按比例绘制。

下面结合附图和实施例，对本发明的具体实施方式作进一步详细描述。

在介绍本申请实施例涉及的技术方案之前，先对部分本申请实施例中的包含的术语进行解释。

(1)计算图(Computational Graph)

本申请实施例中的包含的术语“计算图”，是指一种使用有向图结构对神经网络模型的计算过程进行描述的方式。在实际应用中，有向图结构的基本元素有两个，分别为节点和有向边。其中，节点是指神经网络模型的计算过程中的算子；有向边是指不同算子之间的指向方向，即不同算子之间的依赖关系。

在一些实施例中，神经网络模型可以抽象为张量(tensor)数据和节点，基于张量数据和节点生成该有向图结构。

图1根据本申请的一些实施例，示出了一种计算图10的示意图。如图1所示，计算图10中包括5个节点和5个张量数据。具体地，5个节点分别为节点a、节点b、节点c、节点d和节点e，5个张量数据分别为张量数据T0、张量数据T1、张量数据T2、张量数据T3以及张量数据T4。以节点a、节点b和节点c为例，对节点a来说，节点a与节点b之间存着一条有向边，该有向边的指向方向为节点a指向节点b，该有向边表示节点a生成的张量数据T0为节点b的输入；与此同时，节点a与节点c之间存着一条有向边，该有向边的指向方向为节点a指向节点c，该有向边表示节点a生成的张量数据T1为节点c的输入。

(2)张量

本申请实施例中的包含的术语“张量”，是指张量数据，可以包括神经网络模型中的输入张量数据、输出张量数据，也可以包括特征张量数据等。其中，本申请实施例中的张量数据主要是指各个节点生成的中间数据。

在一些实施例中，神经网络模型可以应用于语音识别、图像识别等场景中。神经网络模型基于输入数据，例如音频数据或者图像数据，可以确定该输入数据在神经网络模型中的张量数据和节点，基于张量数据和节点可以生成对应的计算图。

例如，图2示出本申请的一些实施例的另一种计算图的示意图。如图2所示，计算图20中包括8个节点和7个张量数据。8个节点分别为节点0、节点1、节点2、节点3、节点4、节点5、节点6和节点7，7个张量数据分别为张量数据T0、张量数据T1、张量数据T2、张量数据T3、张量数据T4、张量数据T5、张量数据T6以及张量数据T7。

如图2所示，节点0和节点1之间存在一条有向边，该有向边的指向方向为节点0指向节点1，该有向边表示节点0生成的张量数据T0为节点1的输入；与此同时，节点0与节点3之间存在一条有向边，该有向边的指向方向为节点0指向节点3，该有向边表示节点0生成的张量数据T0为节点3的输入。节点1和节点2之间存在一条有向边，该有向边的指向方向为节点1指向节点2，该有向边表示节点1生成的张量数据T1为节点2的输入。节点2和节点7之间存在一条有向边，该有向边的指向方向为节点2指向节点7，该有向边表示节点2生成的张量数据T2为节点7的输入。节点3和节点4之间存在一条有向边，该有向边的指向方向为节点3指向节点4，该有向边表示节点3生成的张量数据t3为节点4的输入；与此同时，节点3与节点5之间存在一条有向边，该有向边的指向方向为节点3指向节点5，该有向边表示节点3生成的张量数据T3为节点5的输入。节点4和节点6之间存在一条有向边，该有向边的指向方向为节点4指向节点6，该有向边表示节点4生成的张量数据T4为节点6的输入。节点5和节点6之间存在一条有向边，该有向边的指向方向为节点5指向节点6，该有向边表示节点5生成的张量数据T5为节点6的输入。节点6和节点7之间存在一条有向边，该有向边的指向方向为节点6指向节点7，该有向边表示节点6生成的张量数据T6为节点7的输入。

例如，图3示出本申请的一些实施例的另一种计算图的示意图。如图3所示，计算图20中包括7个节点和6个张量数据。7个节点分别为节点0、节点1、节点2、节点3、节点4、节点5和节点6，6个张量数据分别为张量数据T0、张量数据T1、张量数据T2、张量数据T3、张量数据T4、张量数据T5以及张量数据T6。

随着深度学习神经网络越来越大，神经网络模型的张量数据在内存分配过程中可能超出内存的最大容量导致内存分配失败而无法运行。

图4示出本发明实施例提供的神经网络的张量切分方法的流程图，本实施例可适用于对算子进行张量切分的情况，该方法可以由张量切分装置来执行，该装置可以通过软件和/或硬件的方式实现，并一般可集成在人工智能处理器中。如图4所示，所述神经网络的张量切分方法包括但并不限于如下步骤。

在步骤S110中，基于神经网络的输入数据，获取神经网络对应的计算图及张量数据的张量信息。

在本实施例中，神经网络可以为卷积神经网络(convolutional neural network，CNN)、循环神经网络(recurrent neural network，RNN)、深度神经网络(deep neuralnetwork，DNN)等。本申请对神经网络的类型不做任何限制。

在本实施例中，电子设备100获取神经网络对应的计算图，其中，计算图包括N个节点和连接不同节点的有向边，计算图的有向边上承载有张量数据，计算图包括M个张量数据，N为大于1的整数，M为正整数。然后确定各个张量数据的张量名称、张量尺寸以及张量数据之间的依赖关系。可以理解，张量名称可以是张量数据的编号，张量尺寸可以是张量数据的尺寸。

结合图2所示的计算图，可知对应张量数据T0至T6的张量名称可以是T0、T1、T2、T3、T4、T5，张量数据T0的尺寸为4MB，张量数据T1的尺寸为4MB，张量数据T2的尺寸为4MB，张量数据T3的尺寸为4MB，张量数据T4的尺寸为4MB，张量数据T5的尺寸为1MB，张量数据T6的尺寸为4MB。

在步骤S120中，根据计算图的拓扑顺序对神经网络的张量数据进行内存分配，并从所述计算图中获取当前节点和当前张量数据，其中，所述当前张量数据的尺寸超过当前内存的可分配空间。

在本实施例中，节点的输出张量和输入张量之间具有依赖关系，因此相应的张量数据之间也存在依赖关系。例如张量数据T0分配到内存空间后，会等节点1和节点3均完成计算输出张量数据T1以及张量数据T3后，后续节点的计算对张量数据T0不依赖，此时张量数据T0可以从内存中删除以释放内存空间。当前张量数据作为当前节点的输出张量。

根据计算图的拓扑顺序对神经网络的张量数据进行内存分配时，当分配到当前张量数据时，当前张量数据的尺寸超过当前内存的可分配空间，当前张量数据为当前节点的输出张量。

在步骤S130中，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组，其中，所述当前张量切分组包括至少2个张量数据。

在本实施例中，参加图5，步骤S130包括步骤S131至步骤S132。

步骤S131：当神经网络具有分支结构时，从当前节点向前反向搜索至起始节点，所述起始节点为与当前节点相邻的分叉节点。

在本实施例中，分叉节点与多个节点之间存在依赖关系，即分叉节点的输出张量为至少2个节点的输入张量。

步骤S132：从所述起始节点向后搜索至终止节点以获取当前节点分组，所述终止节点为起始节点之后分支的根节点，所述根节点为所有分支的合并节点，当前节点分组包括起始节点、终止节点以及起始节点和终止节点之间的中间节点。

在本实施例中，至少2个分支的输出张量为根节点的输入张量。

步骤S133：将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

在本实施例中，根据张量数据的依赖关系判断当某一张量数据被多个节点作为输入张量，则将该张量数据不需要切分，在划分当前张量切分组时将该张量数据排除。

结合图2，当分配张量数据T5且张量数据T5的尺寸超过当前内存的可分配空间时，从当前节点5反向向前搜索直至节点3作为起始节点，然后从节点3向后搜索至节点6作为终止节点；将节点3和节点6之间的张量数据[T3、T4、T5]作为当前张量切分组。

在一个优选地实施例中，在步骤S132和步骤S133之间还包括以下步骤。

步骤S134：判断当前节点是否为起始节点或者终止节点。

步骤S135：若当前节点为起始节点或者终止节点，从当前节点分组的起始节点继续反向向前搜索以更新当前节点分组的起始节点，更新后的起始节点为当前起始节点相邻的分叉节点以及从所述起始节点向后搜索至终止节点以获取当前节点分组。

当当前节点不为起始节点或终止节点时，执行步骤S136。

在本实施例中，重复步骤S134和步骤S135直至当前节点不为起始节点或终止节点。

结合图2，当分配张量数据T6且张量数据T6的尺寸超过当前内存的可分配空间时，从当前节点6反向向前搜索直至节点3作为起始节点，然后从节点3向后搜索至节点6作为终止节点；由于当前节点6为终止节点，则从起始节点3继续反向向前搜索至相邻的分叉节点0作为更新后的起始节点；然后从起始节点0向后搜索至节点7作为终止节点。将节点0和节点7之间的张量数据[T0、T1、T2、T3、T4、T5、T6]作为当前张量切分组。

步骤S136：遍历当前节点分组内是否存在不支持切分的中间节点。

步骤S137：当所述当前节点分组内存在不支持切分的中间节点，从不支持切分的中间节点处将当前节点分组进行截断以分成至少2个节点分组。

当所述节点分组均为支持切分的中间节点，继续执行步骤S133。

在一个优选地实施例中，步骤S130还包括步骤S138。

步骤S138：当所述当前节点分组的中间节点的个数超过预设阈值时，将所述当前节点分组内去除起始节点与终止节点之间距离近的分支结构以获取缩小后的所述当前节点分组。

在本实施例中，例如起始节点与终止节点之间至少包括2个分支，若所有中间节点的个数超过预设阈值，可以将分支上中间节点数量最少的分支去除以缩小当前节点分组。起始节点与终止节点之间距离近的分支结构指的是起始节点与终止节点之间的中间节点数量最少的分支。

在一个优选地实施例中，当神经网络不具有分支结构时，将当前节点作为当前节点分组以及将当前节点的输入张量和输出张量作为当前张量切分组。

在步骤S140中，对当前张量切分组内的张量数据进行切分以使当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间。

在本实施例中，遍历所有切分策略对当前张量切分组的张量数据进行切分，以及判断当前张量切分组内的切分后的张量数据的尺寸是否超过内存的可分配空间；当当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间时进行内存分配。

应理解，上述实施例中各步骤的序号的大小并不意味着执行顺序的先后，各过程的执行顺序应以其功能和内在逻辑确定，而不应对本申请实施例的实施过程构成任何限定。

本申请的张量切分方法，通过反向搜索获取前节点分组以及当前张量切分组，并对当前张量切分组内的张量数据进行切分以使神经网络的张量数据的读取和存放均在内存中，可以对大的张量数据进行切分以满足内存空间要求，减少对外部存储器的访问，提高计算效率。

进一步地，当所述当前节点分组存在不支持切分的中间节点时，从不支持切分的中间节点处将所述当前节点分组划分成至少2个节点分组，解决节点分组内有些节点的输入输出张量无法切分的问题。

进一步地，当所述节点分组的中间节点超过预设阈值时，对所述节点分组进行缩小处理以获取缩小后的所述节点分组，减少数据重叠和重复计算，进一步提高计算效率。

对应于上文实施例所述的张量切分方法，图6是本申请实施例提供的张量切分装置的结构框图，为了便于说明，仅示出了与本申请实施例相关的部分。如图6所示，所述张量切分装置60包括获取单元61、分配单元62、搜索单元63以及切分单元64。

其中，获取单元61用于基于神经网络的输入数据，获取神经网络对该输入数据进行计算时的计算图以及张量数据的张量信息，其中，所述计算图包括多个计算节点以及各个计算节点之间的张量数据。

分配单元62用于根据计算图的拓扑顺序对神经网络的张量数据进行内存分配，并从所述计算图中获取当前节点和当前张量数据，其中，所述当前张量数据的尺寸超过当前内存的可分配空间。

搜索单元63用于从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组，其中，所述当前张量切分组包括至少2个张量数据。

切分单元64用于对当前张量切分组内的张量数据进行切分以使当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间。

优选地，所述搜索单元63还用于当神经网络具有分支结构时，从当前节点向前反向搜索至起始节点，所述起始节点为与当前节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组，所述终止节点为起始节点之后分支的根节点，所述根节点为所有分支的合并节点，当前节点分组包括起始节点、终止节点以及起始节点和终止节点之间的中间节点；将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

优选地，所述搜索单元63还用于判断当前节点是否为起始节点或者终止节点；若当前节点为起始节点或者终止节点，从当前节点分组的起始节点继续反向向前搜索以更新当前节点分组的起始节点，更新后的起始节点为当前起始节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组；当当前节点不为起始节点或终止节点时，将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组；重复以上步骤直至当前节点不为起始节点或终止节点。

图7是本申请实施例提供的电子设备的结构示意图。如图7所示，该实施例的电子设备70包括：至少一个处理器71(图7中仅示出一个)处理器、存储器72以及存储在所述存储器72中并可在所述至少一个处理器71上运行的计算机程序73，所述处理器71执行所述计算机程序73时实现上述任意各个存储空间分配方法实施例中的步骤。

所述电子设备可以是桌上型计算机、笔记本、掌上电脑及云端服务器等计算设备。该电子设备可包括，但不仅限于，处理器、存储器。本领域技术人员可以理解，图7仅仅是电子设备70的举例，并不构成对电子设备70的限定，可以包括比图示更多或更少的部件，或者组合某些部件，或者不同的部件，例如还可以包括输入输出设备、网络接入设备等。

所称处理器71可以是中央处理单元(Central Processing Unit，CPU)，该处理器71还可以是其他通用处理器、数字信号处理器(Digital Signal Processor，DSP)、专用集成电路(Application Specific Integrated Circuit，ASIC)、现成可编程门阵列(Field－Programmable Gate Array，FPGA)或者其他可编程逻辑器件、分立门或者晶体管逻辑器件、分立硬件组件等。通用处理器可以是微处理器或者该处理器也可以是任何常规的处理器等。

所述存储器72在一些实施例中可以是所述电子设备70的内部存储单元，例如电子设备70的硬盘或内存。所述存储器72在另一些实施例中也可以是所述电子设备70的外部存储设备，例如所述电子设备70上配备的插接式硬盘，智能存储卡(Smart Media Card，SMC)，安全数字(Secure Digital，SD)卡，闪存卡(Flash Card)等。进一步地，所述存储器72还可以既包括所述电子设备70的内部存储单元也包括外部存储设备。所述存储器72用于存储操作系统、应用程序、引导装载程序(Boot Loader)、数据以及其他程序等，例如所述计算机程序的程序代码等。所述存储器72还可以用于暂时地存储已经输出或者将要输出的数据。

本申请实施例还提供了一种计算机可读存储介质，所述计算机可读存储介质存储有计算机程序，所述计算机程序被处理器执行时实现可实现上述各个方法实施例中的步骤。

本申请实施例提供了一种计算机程序产品，当计算机程序产品在电子设备上运行时，使得电子设备执行时实现可实现上述各个方法实施例中的步骤。

所述集成的单元如果以软件功能单元的形式实现并作为独立的产品销售或使用时，可以存储在一个计算机可读取存储介质中。基于这样的理解，本申请实现上述实施例方法中的全部或部分流程，可以通过计算机程序来指令相关的硬件来完成，所述的计算机程序可存储于一计算机可读存储介质中，该计算机程序在被处理器执行时，可实现上述各个方法实施例的步骤。其中，所述计算机程序包括计算机程序代码，所述计算机程序代码可以为源代码形式、对象代码形式、可执行文件或某些中间形式等。所述计算机可读介质至少可以包括：能够将计算机程序代码携带到装置/电子设备的任何实体或装置、记录介质、计算机存储器、只读存储器(ROM，Read－Only Memory)、随机存取存储器(RAM，Random AccessMemory)、电载波信号、电信信号以及软件分发介质。例如U盘、移动硬盘、磁碟或者光盘等。在某些司法管辖区，根据立法和专利实践，计算机可读介质不可以是电载波信号和电信信号。

在上述实施例中，对各个实施例的描述都各有侧重，某个实施例中没有详述或记载的部分，可以参见其它实施例的相关描述。

本领域普通技术人员可以意识到，结合本文中所公开的实施例描述的各示例的单元及算法步骤，能够以电子硬件、或者计算机软件和电子硬件的结合来实现。这些功能究竟以硬件还是软件方式来执行，取决于技术方案的特定应用和设计约束条件。专业技术人员可以对每个特定的应用来使用不同方法来实现所描述的功能，但是这种实现不应认为超出本申请的范围。

在本申请所提供的实施例中，应该理解到，所揭露的装置/电子设备和方法，可以通过其它的方式实现。例如，以上所描述的装置/电子设备实施例仅仅是示意性的，例如，所述模块或单元的划分，仅仅为一种逻辑功能划分，实际实现时可以有另外的划分方式，例如多个单元或组件可以结合或者可以集成到另一个系统，或一些特征可以忽略，或不执行。另一点，所显示或讨论的相互之间的耦合或直接耦合或通讯连接可以是通过一些接口，装置或单元的间接耦合或通讯连接，可以是电性，机械或其它的形式。

所述作为分离部件说明的单元可以是或者也可以不是物理上分开的，作为单元显示的部件可以是或者也可以不是物理单元，即可以位于一个地方，或者也可以分布到多个网络单元上。可以根据实际的需要选择其中的部分或者全部单元来实现本实施例方案的目的。

依照本发明的实施例如上文所述，这些实施例并没有详尽叙述所有的细节，也不限制该发明仅为所述的具体实施例。显然，根据以上描述，可作很多的修改和变化。本说明书选取并具体描述这些实施例，是为了更好地解释本发明的原理和实际应用，从而使所属技术领域技术人员能很好地利用本发明以及在本发明基础上的修改使用。本发明仅受权利要求书及其全部范围和等效物的限制。

Claims (10)

1.一种张量切分方法，其特征在于，包括：

基于神经网络的输入数据，获取神经网络对该输入数据进行计算时的计算图以及张量数据的张量信息，其中，所述计算图包括多个计算节点以及各个计算节点之间的张量数据；

根据计算图的拓扑顺序对神经网络的张量数据进行内存分配，并从所述计算图中获取当前节点和当前张量数据，其中，所述当前张量数据的尺寸超过当前内存的可分配空间；

从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组，其中，所述当前张量切分组包括至少2个张量数据；

对当前张量切分组内的张量数据进行切分以使当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间。

2.根据权利要求1所述的张量切分方法，其特征在于，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组包括：

当神经网络具有分支结构时，从当前节点向前反向搜索至起始节点，所述起始节点为与当前节点相邻的分叉节点；

从所述起始节点向后搜索至终止节点以获取当前节点分组，所述终止节点为起始节点之后分支的根节点，所述根节点为所有分支的合并节点，当前节点分组包括起始节点、终止节点以及起始节点和终止节点之间的中间节点；

将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

3.根据权利要求2所述的张量切分方法，其特征在于，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组还包括：

判断当前节点是否为起始节点或者终止节点；

若当前节点为起始节点或者终止节点，从当前节点分组的起始节点继续反向向前搜索以更新当前节点分组的起始节点，更新后的起始节点为当前起始节点相邻的分叉节点；

从所述起始节点向后搜索至终止节点以获取当前节点分组；

当当前节点不为起始节点或终止节点时，将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组；

重复以上步骤直至当前节点不为起始节点或终止节点。

4.根据权利要求2所述的张量切分方法，其特征在于，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组还包括：

遍历当前节点分组内是否存在不支持切分的中间节点；

所述当前节点分组内存在不支持切分的中间节点，从不支持切分的中间节点处将当前节点分组进行截断以分成至少2个节点分组；

当所述节点分组均为支持切分的中间节点，将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

5.根据权利要求4所述的张量切分方法，其特征在于，从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组还包括：

当所述当前节点分组的中间节点的个数超过预设阈值时，将所述当前节点分组内去除起始节点与终止节点之间距离近的分支结构以获取缩小后的所述当前节点分组。

6.根据权利要求1所述的张量切分方法，其特征在于，对当前张量切分组内的张量数据进行切分包括：

遍历所有切分策略对当前张量切分组的张量数据进行切分；

判断当前张量切分组内的切分后的张量数据的尺寸是否超过内存的可分配空间；

当当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间时进行内存分配。

7.一种张量切分装置，其特征在于，包括：

获取单元，用于基于神经网络的输入数据，获取神经网络对该输入数据进行计算时的计算图以及张量数据的张量信息，其中，所述计算图包括多个计算节点以及各个计算节点之间的张量数据；

分配单元，用于根据计算图的拓扑顺序对神经网络的张量数据进行内存分配，并从所述计算图中获取当前节点和当前张量数据，其中，所述当前张量数据的尺寸超过当前内存的可分配空间；

搜索单元，用于从当前节点进行反向搜索以获取当前节点分组以及当前张量切分组，其中，所述当前张量切分组包括至少2个张量数据；

切分单元，用于对当前张量切分组内的张量数据进行切分以使当前张量切分组内的切分后的张量数据的尺寸均不超过当前内存的可分配空间；

其中，所述搜索单元还用于当神经网络具有分支结构时，从当前节点向前反向搜索至起始节点，所述起始节点为与当前节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组，所述终止节点为起始节点之后分支的根节点，所述根节点为所有分支的合并节点，当前节点分组包括起始节点、终止节点以及起始节点和终止节点之间的中间节点；将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组。

8.根据权利要求7所述的张量切分装置，其特征在于，所述搜索单元还用于判断当前节点是否为起始节点或者终止节点；若当前节点为起始节点或者终止节点，从当前节点分组的起始节点继续反向向前搜索以更新当前节点分组的起始节点，更新后的起始节点为当前起始节点相邻的分叉节点；从所述起始节点向后搜索至终止节点以获取当前节点分组；当当前节点不为起始节点或终止节点时，将起始节点和终止节点之间的张量数据以及终止节点的输出张量划分为当前张量切分组；重复以上步骤直至当前节点不为起始节点或终止节点。

9.一种电子设备，包括存储器、处理器以及存储在所述存储器中并可在所述处理器上运行的计算机程序，其特征在于，所述处理器执行所述计算机程序时实现如权利要求1至6任一项所述的方法。

10.一种计算机可读存储介质，所述计算机可读存储介质存储有计算机程序，其特征在于，所述计算机程序被处理器执行时实现如权利要求1至6任一项所述的方法。