一种神经网络计算图的分割方法、装置、设备及存储介质与流程

本技术涉及神经网络，尤其涉及一种神经网络计算图的分割方法、装置、设备及存储介质。

背景技术：

1、人工智能近年得到迅速发展，在图像分类、检测、视频和语音处理等领域取得良好的应用效果，并且依然具备广阔的发展前景。神经网络是人工智能应用的核心，深度学习神经网络算法是其中最常见的一种神经网络模型。深度学习神经网络的工作负载特征是计算和数据密集。训练神经网络所需的训练以及后续推理所需参数则通常为m至gb字节数量级。由于常规计算机，甚至单个专用服务器无法提供如此高的算力并配备足够内存，因此越来越多的神经网络训练和推理在专门的神经网络平台上被分布式执行。

2、分布式执行需要对神经网络的计算图进行切分，并将切分后的子图部署到合适的处理设备上，以提高神经网络的推理速度。目前针对神经网络计算图的切分方案只适合简单的算子链式相连的神经网络计算图，而越来越多的神经网络计算图包含更复杂的算子连接依赖关系，这些都涉及非链式的神经网络计算图。因此，如何对任意的神经网络计算图进行子图切分成为亟待解决的问题。

技术实现思路

1、本技术提供一种神经网络计算图的分割方法、装置、设备及存储介质，能够实现对任意的神经网络计算图进行子图切分。

2、为达到上述目的，本技术采用如下技术方案：

3、本技术实施例第一方面，提供了一种神经网络计算图的分割方法，该方法包括：

4、获取待分割的计算图，所述计算图中包括多个节点以及多个节点之间的边，所述节点中包括计算单元，所述边用于指示计算间的依赖关系；

5、根据各所述节点的处理类型为所述计算图中的各个节点标记对应的设备标识，得到标记后的第一计算图，所述设备标识用于指示各所述节点的执行设备；

6、针对所述第一计算图中具有相同设备标识的节点，根据预定义的子网络结构模板以及各相邻节点之间的融合代价进行节点融合，得到融合后的第二计算图；

7、针对所述第二计算图，根据各所述执行设备对应的内存对所述第二计算图进行分割，得到分割后的目标计算图，所述目标计算图中包括多个子计算图；

8、根据各个子计算图对应的设备标识，将各个子计算图部署至对应的分布式设备上进行并行处理。

9、在一种可能的实现方式中，所述根据各所述节点的处理类型为所述计算图中的各个节点标记对应的设备标识，包括：

10、获取当前计算场景对应的执行设备，以及各所述执行设备对应的设备标识，根据各所述节点的处理类型为各所述节点匹配对应的执行设备，并将各所述执行设备的设备标识标记至对应的节点。

11、在一种可能的实现方式中，所述根据预定义的子网络结构模板以及各相邻节点之间的融合代价进行节点融合，得到融合后的第二计算图，包括：

12、遍历所述第一计算图，针对所述第一计算图中具有相同设备标识的节点，将符合所述子网络结构模板的节点进行融合，得到融合后的中间计算图；

13、针对所述中间计算图中具有相同设备标识的节点，根据各相邻节点之间的融合代价进行节点融合，得到融合后的第二计算图。

14、在一种可能的实现方式中，所述针对所述第一计算图中具有相同设备标识的节点，将符合所述子网络模板的节点进行融合之前，所述方法包括：

15、根据当前计算场景，预定义多个子网络结构，并将配置的多个子网络结构作为所述子网络结构模板。

16、在一种可能的实现方式中，所述根据各相邻节点之间的融合代价进行节点融合，得到融合后的第二计算图，包括：

17、针对所述中间计算图中具有相同设备标识的节点，获取各所述相邻节点之间的数据传输量和通信带宽；

18、利用贪心算法根据各所述相邻节点之间的数据传输量和通信带宽，确定出各相邻节点之间的融合代价，以及根据各所述融合代价确定出融合路径，所述融合路径包括待融合的节点和边；

19、根据所述融合路径对所述中间计算图进行融合，得到融合后的第二计算图。

20、在一种可能的实现方式中，所述针对所述第二计算图，根据各所述执行设备对应的内存对所述第二计算图进行分割，包括：

21、获取各所述执行设备对应的运行内存，并获取所述第二计算图中的各子计算图所需的处理内存；

22、遍历所述第二计算图，若所述第二计算图中的子计算图所需的处理内存大于对应的执行设备的运行内存，则对所述第二计算图中的子计算图进行分割。

23、在一种可能的实现方式中，所述将各个子计算图部署至对应的设备上进行处理，包括：

24、获取所述目标计算图中可并行处理的目标子计算图，为所述目标子计算图开设多个线程来并行处理。

25、本技术实施例第二方面，提供了一种神经网络计算图的分割装置，该装置包括：

26、获取模块，用于获取待分割的计算图，所述计算图中包括多个节点以及多个节点之间的边，所述节点中包括计算单元，所述边用于指示计算间的依赖；

27、处理模块，用于根据各所述节点的处理类型为所述计算图中的各个节点标记对应的设备标识，得到标记后的第一计算图，所述设备标识用于指示各所述节点的执行设备；

28、融合模块，用于针对所述第一计算图中具有相同设备标识的节点，根据预定义的子网络结构模板以及各相邻节点之间的融合代价进行节点融合，得到融合后的第二计算图；

29、分割模块，用于针对所述第二计算图，根据各所述执行设备对应的内存对所述第二计算图进行分割，得到分割后的目标计算图，所述目标计算图中包括多个子计算图；

30、部署模块，用于根据各个子计算图对应的设备标识，将各个子计算图部署至对应的分布式设备上进行并行处理。

31、本技术实施例第三方面，提供了一种电子设备，包括存储器和处理器，所述存储器存储有计算机程序，所述计算机程序被所述处理器执行时实现本技术实施例第一方面中的神经网络计算图的分割方法。

32、本技术实施例第四方面，提供了一种计算机可读存储介质，其上存储有计算机程序，所述计算机程序被处理器执行时实现本技术实施例第一方面中的神经网络计算图的分割方法。

33、本技术实施例提供的技术方案带来的有益效果至少包括：

34、本技术实施例提供的神经网络计算图的分割方法，通过获取待分割的计算图，该计算图中包括多个节点以及多个节点之间的边，其中，节点中包括计算单元，边用于指示计算间的依赖，然后根据各所述节点的处理类型为所述计算图中的各个节点标记对应的设备标识，得到标记后的第一计算图，所述设备标识用于指示各所述节点的执行设备，再针对所述第一计算图中具有相同设备标识的节点，根据预定义的子网络结构模板以及各相邻节点之间的融合代价进行节点融合，得到融合后的第二计算图，以及针对所述第二计算图，根据各所述执行设备对应的内存对所述第二计算图进行分割，得到分割后的目标计算图，所述目标计算图中包括多个子计算图，最后根据各个子计算图对应的设备标识，将各个子计算图部署至对应的分布式设备上进行并行处理，这样可以实现对任意的神经网络计算图进行子图切分，进而可以提高神经网络的处理效率。