[发明专利]一种基于深度强化学习的流水线并行训练任务分配方法

权利要求书

1.一种基于深度强化学习的流水线并行训练任务分配方法，其特征在于，包括如下步骤：

(1)初始化深度学习模型，并把该模型载入；

(2)建立待训练模型短暂运行的分析方案：记录模型详细信息，包括模型总层数、每一层的名称、计算时间、计算量、参数量、激活值大小，利用记录的模型信息构建待训练模型参数文本库；

(3)根据获取的模型参数文本信息，构建任务分配预测网络；

(4)使用任务分配预测网络为待训练模型生成流水线并行训练方案；

(5)按照该流水线并行训练方案执行训练，采样求平均得到单个批次平均训练时间；

(6)计算任务分配方案的奖励函数，反向更新任务分配预测网络中的参数；

(7)判断是否生成训练时间最小的任务分配方案，是则输出该方案，否则返回步骤(4)继续迭代；

(8)按照输出方案将模型部署到异构计算节点中，得到针对拟训练目标网络的流水线并行训练分配方案。

2.根据权利要求1所述的基于深度强化学习的流水线并行训练任务分配方法，其特征在于，所述步骤(3)中任务分配预测网络结构如下：

(3.1)任务划分网络，负责将模型中所有层划分为若干分组；

(3.2)任务分配网络，负责将每个分组任务分配到对应的计算节点中。

3.根据权利要求1所述的基于深度强化学习的流水线并行训练任务分配方法，其特征在于，所述步骤(4)中流水线并行训练方案的具体步骤如下：

(4.1)基于前馈神经网络的任务划分；

(4.2)基于注意力机制的任务分配。

4.根据权利要求3所述的基于深度强化学习的流水线并行训练任务分配方法，其特征在于，所述步骤(4.1)中基于前馈神经网络的任务划分的具体步骤如下：

(4.1.1)在模型中，每一层都对应着不同的操作，将其转换为词向量；

(4.1.2)模型每一层的参数指标包括参数量、计算量与生成激活值大小；根据模型结构定义计算每一层需要的参数量；根据每一层参数量与对应计算操作估算其需要的计算量；根据每一层参数量与对应计算操作，估算出其每一层的生成的激活值大小；

(4.1.3)对模型逐层编号；

(4.1.4)遍历模型中的每一层，经过上述的处理步骤后，每一层得到向量信息，将这些信息输入到任务划分网络，预测出DNN模型每一层的划分结果——模型分组结果。

5.根据权利要求3所述的基于深度强化学习的流水线并行训练任务分配方法，其特征在于，所述步骤(4.2)中基于注意力机制的任务分配的具体步骤如下：

(4.2.1)将任务划分网络预测得到的模型分组结果，通过编码器处理生成语义向量，顺序传入解码器；

(4.2.2)解码器按照分组顺序分配计算节点；

(4.2.3)在解码器进行预测的每个步骤中利用注意力机制，将上一步骤所分配的计算节点编号传入解码器，生成当前步骤计算节点分配结果；

(4.2.4)将模型划分结果与任务分配结果部署在实际计算节点中，得到整体的任务分配方案。