一种基于强化学习的计算图自动划分方法

1.本发明属于资源分配和强化学习领域，涉及一种使用强化学习算法解决深度学习计算图划分问题的方法。


背景技术：

2.近年来，深度学习已经在图像分析、自然语言处理、语音识别、视频分类等领域取得了令人瞩目的成功。但是，深度学习依赖于强大的计算能力，优化深度学习的系统框架以降低算力需求对于深度学习的应用有重要的作用。面对深度学习模型爆发式的计算需求，针对ai领域的ai芯片被广泛应用。
3.ai芯片普遍采用了众核架构。ai芯片用于专门处理ai应用中的大量计算任务，而其他非计算任务仍由cpu处理。ai芯片集成了多个内核，功能上可以具有针对性地加速某类特定的算法或任务。市面上对ai芯片的定义更多来自于功能方面，而不拘泥于它的架构。近年来，与深度学习相关的ai芯片产品相继出现，从谷歌、英特尔、英伟达等科技巨头到寒武纪、地平线等创业公司，都推出了自己的解决方案。随着芯片设计技术的发展成熟，ai芯片架构随着ai技术的发展实现了快速迭代。
4.深度学习编译器往往将不同框架实现的模型拆分为多个子任务部署在众核芯片上，使用流水线结构处理计算任务，以达到最优性能。在流水线结构中，一个大的计算任务可以分为若干个并行的子任务集，其中每一个子任务集由众核芯片中的几个核并行处理，子任务集到核资源集的映射由运行时系统(run-time system)完成。当运行某个深度学习模型作为计算任务时，划分子任务即为对这个深度学习模型的计算图划分为多个子图。在计算图中，用节点模拟计算过程，优点是可以分割运算，甚至采用多gpu的方式进行运算。
5.在流水线结构中如何以一种优化的方式把计算任务分配到各个处理器是非常重要的，这个问题通常称为负载平衡问题。负载平衡是并行计算中的一个基本问题，通过最大限度地减少处理器的空闲时间和处理器之间的通信时间，达到使并行应用性能最大化的效果。
6.流水线结构的调度可以减少处理器的空闲时间，提高程序执行性能，增大硬件资源的利用率。然而，如处理器核性能、片上存储、通信带宽等系统资源的的限制会影响软件流水调度性能。
7.强化学习在解决资源调度方面，已经取得了一定的成果。2017年mirhoseini等人提出使用强化学习来优化tensorflow模型在分布式系统中的计算图节点分配问题。该论文使用了一个序列到序列(seq2seq)的模型。模型由编码器和解码器组成，计算图节点按拓扑排序输入到模型，各个节点被分配的设备作为输出。此后许多基于强化学习的资源调度方案被提出。addanki等人使用强化学习算法来实现神经网络在分布式资源上的调度。此研究不断迭代资源分配方案，而非一次性得到计算图的节点分配方案。而后的研究也普遍利用强化学习来解决分布式系统中的任务调度问题，而差异大多体现在具体的深度学习模型上。但是以上用强化学习解决资源调度问题的方法均是按照资源布局分配资源，不适应处
理器的核资源布局改变的情况。luo针对分布式流处理系统的资源调度问题，提出了基于深度强化学习和多级图划分思想的算法，按照资源数量分配资源。通过图粗化降低图形复杂度，将大规模数据流图处理为小规模数据流图后，利用强化学习进行训练，再将结果映射到大规模数据流图中。但是此方法对数据流图进行粗化处理后，限定了强化学习寻找最优解的空间，使得强化学习的效果收到了一定的限制。
8.本发明主要研究在众核芯片架构上训练深度学习模型时，如何有效划分模型对应的计算图并实现负载均衡。故本文将要设计一种自动划分深度学习计算图并为每一个子图分配核资源数量的算法，使得深度学习模型在众核芯片运行时间最少，是一种按照资源数量分配资源的方法。


技术实现要素：

9.本发明的目的，在于提供了一种基于强化学习的计算图自动划分方法，该方法能够自动将深度学习模型对应的深度学习计算图划分为子图，并且按照资源数量为每一个子图分配核资源，达到缩短深度学习模型运行时间的目标。
10.为达到上述目的，本发明所述的基于强化学习的计算图自动划分方法包括以下步骤：
11.步骤1，对计算图进行拓扑排序转化为线性表；
12.步骤2，将待划分的计算图和众核处理器建模为强化学习中的马尔可夫决策过程，提取当前计算图的子图划分情况与核资源分配情况作为强化学习中的状态，调整相邻两个核之间的层数分配作为强化学习中的动作，计算图在众核处理器上的运行时间和存储情况作为强化学习中的奖励；
13.步骤3，利用reinforce算法对马尔可夫决策过程求解，训练图划分算法模型。
14.上述步骤1的具体过程是:
15.步骤11,对深度学习计算图进行拓扑排序，将图转化成线性表的结构。线性表中的元素排列顺序与节点的运算顺序相一致，线性表中的数据元素与深度学习模型中的层相对应；
16.步骤12,设法表示出当前每个层的数据操作类型与操作的超参数、图中各边的数据量大小，从而计算出计算图中的总节点数和每个操作的运算量、需要的存储量以及路由量。
17.上述步骤2具体包括如下步骤:
18.步骤21,提取当前计算图的子图划分情况与核资源分配情况作为强化学习中的状态：状态由两部分组成，第一部分为计算图的节点划分和资源分配状态，第二部分是各子图的运算量状态；
19.步骤22,将调整相邻两个核之间的层数分配作为强化学习中的动作，分别有四类动作：将两个相邻核的所有层合并到后一个核上处理、将两个相邻核的所有层合并到前一个核上处理、将前一个核处理的最后一个层(即线性表中的最后一个层)交给后一个核处理、将后一个核处理的第一个层(即线性表中的第一个层)交给前一个核处理；
20.步骤23,将计算图在众核芯片上的运行时间和存储情况作为强化学习中的奖励：奖励值设定为reward＝-a*max(t) b，其中t＝{t1,t2,
…
,tk}为计算图g在众核处理器m上的
运行时间情况，s＝{s1,s2,
…
,sk}表示众核处理器m上的数据存储情况。若max(s)超过限制，则赋予b一个惩罚值。训练的目标是令奖励值尽可能大。
21.上述步骤3具体包括如下步骤:
22.步骤31,初始化整个图划分环境，导入转化为线性表结构的深度学习计算图，统计节点总数，初始化核资源分配情况，并根据核资源总数初始化可供选择的动作，将记录子图内的核数量的变量、记录各个核资源状态的变量、记录划分节点数的变量、记录奖励值的变量等全部变为初始状态，并初始化所有动作的概率；
23.步骤32,根据动作概率选择动作a，执行动作a后当前环境状态s改变；
24.步骤33,根据当前的状态计算每一个子图的计算量、存储量和路由量，将每个子图中的这三个值分别进行比较，得出情况最差的三个值进行加权运算计算出奖励值r；
25.步骤34,判断奖励值是否达到预设要求：如果达到则更新动作概率且本回合结束，如果未达到则继续选择动作，重复步骤32-步骤34；
26.步骤35,将本回合所有的s,a,r输入到神经网络中进行训练，更新选择动作的概率；
27.步骤36,达到设定的回合数后，训练过程结束，保存神经网络模型。
28.以上实施例仅用以说明本发明的技术方案,而非对其限制，凡是按照本发明提出的技术思想,在技术方案基础上所做的任何改动,均落入本发明保护范围之内。
附图说明
29.图1为基于强化学习的计算图自动划分方法流程图；
30.图2为基于深度强化学习的面向众核芯片的计算图自动划分方法结构的示意图；
31.图3为神经网络结构图。
具体实施方式
32.下面结合附图对本发明做进一步详细描述：
33.本发明提供一种基于强化学习的计算图自动划分方法,包含三个部分，分别是对计算图进行拓扑排序、将计算图划分问题建模成强化学习的马尔可夫决策过程，利用基于reinforce算法的面向众核处理器的计算图自动划分算法训练神经网络模型。具体实现方法如下:
34.定义面向众核处理器的计算图划分问题为：对于一个计算图，对于一个计算图g＝(o,e)，o＝{op1,op2,
···
opm}是图上算子的集合，e是e条边的集合。将计算图g划分为子图集g’＝{g1,g2,
···
,gk}，其中k≤m，即每个子图g由一个或多个算子组成。计算图g部署在众核处理器m上运行，m上集成了n个核资源。为每个子图分配核数，c＝{c1,c2,
…
,ck}为核资源组的集合，表示为子图gi分配了ci个核资源，其中c1 c2
···
ck＝n，负责子图gi的ci个核资源构成了一个核资源组。每个划分方案p＝(g’,c)表示将计算图g以g’的方式划分子图，并且以c的核数分配方案为每个子图分配核资源。在划分方案p下，s(p)＝{s1,s2,
…
,sk}表示众核处理器m上的数据存储情况，其中si表示第i个核资源组中的单核存储量，t(p)＝{t1,t2,
…
,tk}为计算图g在众核处理器m上的运行时间情况，其中ti表示子图gi的运行时间。训练目标是要寻找一个分配方案p＝(g’,c)使得max(t)最短，且max(s)不超过众核处理器m
的存储上限。
35.根据上述分析，本发明涉及了基于强化学习的计算图自动划分方法，步骤如下：
36.步骤a，对计算图进行拓扑排序转化为线性表；
37.步骤b，将待划分的计算图和众核处理器建模为强化学习中的马尔可夫决策过程；
38.步骤c，利用reinforce算法对马尔可夫决策过程求解，训练图划分算法模型。
39.进一步的，步骤a的描述详见发明内容。
40.步骤b具体为：
41.步骤b1，将状态建模为当前计算图的子图划分情况与核资源分配情况。由于影响负载均衡的主要因素是子图的运算量，故状态由两部分组成，第一部分为计算图的节点划分和资源分配状态，第二部分是各子图的运算量状态。
42.计算图的节点划分和资源分配状态用列表arranged_core来记录，表示众核处理器上各核资源处理计算图的情况。处理器包含n个核资源{core1,core2,
···
,coren}，列表arranged_core表示为[layers1,layers2,layers3,layers4,
…
,layersn]，其中且0≤layersi≤m，可将layersi看作corei处理的层数。若layersi均不为0，则整张计算图被划分为n个子图，拥有的层数分别为layers1,layers2,layers3,layers4，
…
，layersn。当出现layersi，layers
i 1
，
…
，layers
i p
均为0且layers
i p 1
≠0，即arranged_core中连续出现p 1个0时(0≤p≤m-1)，则从第i个核资源起，直到第i p个核资源，共同组成一个核资源组，处理layers
i p 1
个层构成的一个子图。需要强调的是，在实际运行过程中，corei，core
i 1
，
…
，core
i p
并不是处理0个层，此处的状态设置仅为方便强化学习的建模。
[0043]
各子图的运算量状态用列表arranged_macs来记录，表示arranged_core列表中layersi个层的计算量。arranged_macs与arranged_core的设定相对应，表示为[macs1,macs2,macs3,macs4…
macsn]，其中macsi为layersi个层的总运算量。关于macsi为0的定义与layersi相类似。
[0044]
步骤b2，在当前的划分状态下，调整相邻两个核之间的层数分配建模为一次动作。
[0045]
对相邻的两个核调整层数的动作共有四类。对相邻核corei,core
i 1
可以采取的四类动作分别为：(a)将两个相邻核corei,core
i 1
上的层合并到后一个核上core
i 1
上处理。即corei合并到core
i 1
的核资源组中，新核资源组需要承担layersi的任务。(b)将两个相邻核corei,core
i 1
上的层合并到前一个核上corei上处理。即corei的核资源组处理layersi,layers
i 1
中的所有层，并且core
i 1
加入core
i 2
的核资源组。(c)将前一个核资源corei处理的layersi个层中最后一个层(即线性表中的最后一个层)交给后一个核资源core
i 1
处理。(d)将后一个核资源core
i 1
处理的layers
i 1
个层中第一个层(即线性表中的第一个层)交给前一个核资源corei处理。
[0046]
动作空间为所有可选择的动作集合。本文对相邻两个核采取动作，对于n个核，共有n-1组相邻核，每类动作有n-1个选择。一共有四类动作，所以动作空间大小为4*(n-1)。表1展示了这四类动作的定义，前两类动作为合并动作，后两类动作为拆分动作。为了便于代码编写，动作编号始于0。
[0047]
表1四类action的定义
[0048][0049]
在执行动作时，可能会发生该动作对此时环境状态无效的情况，那么令环境状态保持不变。比如两个核原本处于同一核资源组中，此时层数的调整无效果，则动作不会改变状态。
[0050]
步骤b3，将计算图在众核芯片上的运行时间和存储情况建模为奖励。本文的目标是要寻找一个分配方案p＝(g’,c)使得max(t)最短，且max(s)不超过众核处理器m的存储上限。故奖励值设定为reward＝-a*max(t) b，若max(s)超过限制，则赋予b一个惩罚值。训练的目标是令奖励值尽可能大。
[0051]
进一步的，步骤c的训练图划分算法模型详细步骤如下：
[0052]
步骤c1,初始化整个图划分环境。导入转化为线性表结构的深度学习计算图，统计节点总数，初始化核资源分配情况，并根据核资源总数初始化可供选择的动作，将记录子图内的核数量的变量、记录各个核资源状态的变量、记录划分节点数的变量、记录奖励值的变量等全部变为初始状态，并初始化所有动作的概率，记总核数为n,令m＝n-1；
[0053]
步骤c2，根据动作概率选择动作a；
[0054]
步骤c3，根据a的值选择相应的动作，执行动作a；
[0055]
步骤c3,根据当前的状态s计算每一个子图的计算量、存储量和路由量，将每个子图中的这三个值分别进行比较，得出情况最差的三个值进行加权运算计算出奖励值r；
[0056]
步骤c4,判断r是否达到预设要求：如果达到则更新动作概率且本回合结束，如果未达到则继续选择动作，重复步骤32-步骤34；
[0057]
步骤c5,将本回合所有的s,a,r输入到神经网络中进行训练，更新选择动作的概率(图3展示了神经网络的结构)；
[0058]
步骤c6,达到设定的回合数后，训练过程结束，保存神经网络模型。
[0059]
以上实施例仅用以说明本发明的技术方案,而非对其限制，凡是按照本发明提出的技术思想,在技术方案基础上所做的任何改动,均落入本发明保护范围之内。