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

本发明专利技术公开了一种基于强化学习的计算图自动划分方法，包括以下步骤：步骤1，对计算图进行拓扑排序转化为线性表；步骤2，将待划分的计算图和众核处理器建模为强化学习中的马尔可夫决策过程，提取当前计算图的子图划分情况与核资源分配情况作为强化学习中的状态，调整相邻两个核之间的层数分配作为强化学习中的动作，计算图在众核处理器上的运行时间和存储情况作为强化学习中的奖励；步骤3，利用REINFORCE算法对马尔可夫决策过程求解，训练图划分算法模型。图划分算法模型。图划分算法模型。

【技术实现步骤摘要】
一种基于强化学习的计算图自动划分方法


[0001]本专利技术属于资源分配和强化学习领域，涉及一种使用强化学习算法解决深度学习计算图划分问题的方法。

技术介绍

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

技术实现思路

[0009]本专利技术的目的，在于提供了一种基于强化学习的计算图自动划分方法，该方法能够自动将深度学习模型对应的深度学习计算图划分为子图，并且按照资源数量为每一个子图分配核资源，达到缩短深度学习模型运行时间的目标。
[0010]为达到上述目的，本专利技术所述的基于强化学习的计算图自动划分方法包括以下步骤：
[0011]步骤1，对计算图进行拓扑排序转化为线性表；
[0012]步骤2，将待划分的计算图和众核处理器建模为强化学习中的马尔可夫决策过程，提取当前计算图的子图划分情况与核资源分配情况作为强化学习中的状态，调整相邻两个核之间的层数分配作为强化学习中的动作，计算图在众核处理器上的运行时间和存储情况作为强化学习中的奖励；
[0013]步骤3，利用REINFORCE算法对马尔可夫决策过程求解，训练图划分算法模型。
[0014]上述步骤1的具体过程是:
[0015]步骤11,对深度学习计算图进行拓扑排序，将图转化成线性表的结构。线性表中的元素排列顺序与节点的运算顺序相一致，线性表中的数据元素与深度学习模型中的层相对应；
[0016]步骤12,设法表示出当前每个层的数据操作类型与操作的超参数、图中各边的数据量大小，从而计算出计算图中的总节点数和每个操作的运算量、需要的存储量以及路由量。
[0017]上述步骤2具体包括如下步骤:
[0018]步骤21,提取当前计算图的子图划分情况与核资源分配情况作为强化学习中的状态：状态由两部分组成，第一部分为计算图的节点划分和资源分配状态，第二部分是各子图的运算量状态；
[0019]步骤22,将调整相邻两个核之间的层数分配作为强化学习中的动作，分别有四类动作：将两个相邻核的所有层合并到后一个核上处理、将两个相邻核的所有层合并到前一个核上处理、将前一个核处理的最后一个层(即线性表中的最后一个层)交给后一个核处理、将后一个核处理的第一个层(即线性表中的第一个层)交给前一个核处理；
[0020]步骤23,将计算图在众核芯片上的运行时间和存储情况作为强化学习中的奖励：奖励值设定为reward＝
‑
a*max(t)+b，其中T＝{t1,t2,
…
,t
k
}为计算图G在众核处理器M上的
运行时间情况，S＝{s1,s2,
…
,s
k
}表示众核处理器M上的数据存储情况。若max(s)超过限制，则赋予b一个惩罚值。训练的目标是令奖励值尽可能大。
[0021]上述步骤3具体包括如下步骤:
[0022]步骤31,初始化整个图划分环境，导入转化为线性表结构的深度学习计算图，统计节点总数，初始化核资源分配情况，并根据核资源总数初始化可供选择的动作，将记录子图内的核数量的变量、记录各个核资源状态的变量、记录划分节点数的变量、记录奖励值的变量等全部变为初始状态，并初始化所有动作的概率；
[0023]步骤32,根据动作概率选择动作a，执行动作a后当前环境状态s改变；
[0024]步骤33,根据当前的状态计算每一个子图的计算量、存储量和路由量，将每个子图中的这三个值分别进行比较，得出情况最差的三个值进行加权运本文档来自技高网...

【技术保护点】

【技术特征摘要】
1.一种基于强化学习的计算图自动划分方法，其特征在于，包括以下步骤：步骤1，对计算图进行拓扑排序转化为线性表；步骤2，将待划分的计算图和众核处理器建模为强化学习中的马尔可夫决策过程，提取当前计算图的子图划分情况与核资源分配情况作为强化学习中的状态，调整相邻两个核之间的层数分配作为强化学习中的动作，计算图在众核处理器上的运行时间和存储情况作为强化学习中的奖励；步骤3，利用REINFORCE算法对马尔可夫决策过程求解，训练图划分算法模型。2.如权利要求1所述的基于强化学习的计算图自动划分方法,其特征在于:所述步骤1的具体过程是:步骤11,对深度学习计算图进行拓扑排序，将图转化成线性表的结构。线性表中的元素排列顺序与节点的运算顺序相一致，线性表中的数据元素与深度学习模型中的层相对应；步骤12,设法表示出当前每个层的数据操作类型与操作的超参数、图中各边的数据量大小，从而计算出计算图中的总节点数和每个操作的运算量、需要的存储量以及路由量。3.根据权利要求1所述的基于强化学习的计算图自动划分方法，其特征在于，所述步骤2的具体过程是:步骤21,提取当前计算图的子图划分情况与核资源分配情况作为强化学习中的状态：状态由两部分组成，第一部分为计算图的节点划分和资源分配状态，第二部分是各子图的运算量状态；步骤22,将调整相邻两个核之间的层数分配作为强化学习中的动作，分别有四类动作：将两个相邻核的所有层合并到后一个核上处理、将两个相邻核的所有层合并到前一个核上处理、将前一个核处理的最后一个层(即线性表中的最后一个层)交给后一个核处理、将后一个核处理的第一个层(即线性表中的第一个层)交...

【专利技术属性】
技术研发人员：崔毅东，林孟群，雷友珣，陈莉萍，
申请(专利权)人：北京邮电大学，
类型：发明
国别省市：