[发明专利]一种基于强化学习内存调度决策的模型训练方法及系统

权利要求书

1.一种基于强化学习内存调度决策的模型训练方法，其特征在于，包括以下步骤：

步骤一：在第一轮迭代训练过程中，初始化调度方案；

步骤二：采集上一轮次神经网络模型迭代训练完后的GPU的内存空间、转移带宽，以及神经网络模型的网络层信息，记录数据的大小、各层的依赖关系；并根据采集的信息计算数据交换时间和训练时间；

步骤三：基于步骤二获取的信息动态决策每轮次的调度方案，包括以下子步骤：

(3.1)对步骤二获得的训练时间、内存空间消耗进行归一化处理后，进行加权评估作为上一轮次的奖励收益；

(3.2)根据步骤(3.1)获得的上一轮次的奖励收益和上一轮次获得的反馈值Q采用ε-greedy算法，选择当前轮次等待交换的张量和当前轮次的调度方案，神经网络模型依据调度方案进行当前轮次训练；

(3.3)根据上一轮次的反馈值Q、奖励收益以及当前轮次的调度方案，对反馈值Q进行更新并存储至Q表；

步骤四：重复步骤二、三进行迭代训练；随着迭代次数增加，调度方案将收敛到调度最优解；当达到预期训练效果时，结束训练。

2.根据权利要求1所述的模型训练方法，其特征在于，所述步骤(3.1)中，上一轮次的奖励收益R_plan具体如下：

R_plan＝R_time+w(R_mem-R_time)

其中，R_time表示上一轮次的训练时间，R_mem表示上一轮次的训练消耗的GPU内存空间，w表示在奖励收益R_plan中，上一轮次的训练消耗的GPU内存空间R_mem所占据的比例。

3.根据权利要求1所述的模型训练方法，其特征在于，所述步骤(3.3)中，对反馈值Q进行更新具体为：

Q_k(s,t)＝Q_k-1(s,t)+α[R_plan+γQ_k-1(s',t')-Q_k-1(s,t)]

其中，k表示当前轮次，k-1表示上一轮次，t是上一轮次被交换的数据张量，s是上一轮次采用的调度方案，t'、s'分别是当前轮次被交换的数据张量和采用的调度方案，R_plan是上一轮次的奖励收益，α是学习率，γ是折扣因子；Q_k-1(s,t)表示第k-1轮次，被交换的数据张量为t和采用的调度方案为s时的反馈值，Q_k-1(s',t')表示第k-1轮次，被交换的数据张量为t'和采用的调度方案为s'时的反馈值，Q_k-1(s,t)、Q_k-1(s',t')依据Q表查找获得。

4.根据权利要求1所述的模型训练方法，其特征在于，所述步骤(3.2)中，采用ε-greedy算法，选择当前轮次等待交换的张量和当前轮次的调度方案时，每轮次只选择同一种网络层类型对应的张量作为等待交换的张量。

5.根据权利要求1所述的模型训练方法，其特征在于，所述模型训练方法包括预训练和正式训练，其中，预训练采用部分训练数据进行训练并动态决策每轮次的调度方案，正式训练采用全部训练数据在预训练的基础上继续进行训练并动态决策每轮次的调度方案。

6.一种基于强化学习内存调度决策的模型训练系统，其特征在于，用于实现权利要求1-5任一项基于强化学习内存调度决策的模型训练方法，包括：

采集模块，用于采集上一轮次神经网络模型迭代训练完后的GPU的内存空间、转移带宽，以及神经网络模型的网络层信息，记录数据的大小、各层的依赖关系；并根据采集的信息计算数据交换的对应时间和训练时间；

决策模块，用于基于采集模块获取的信息动态决策每轮次的调度方案；

训练模块，依据采集模块和决策模块的数据进行迭代训练；随着迭代次数增加，调度方案将收敛到调度最优解；当达到预期训练效果时，结束训练。