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

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

背景技术：

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

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

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

技术实现思路

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

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

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

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

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

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

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

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

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

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

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

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

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

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