一种基于强化学习内存调度决策的模型训练方法及系统技术方案

本发明专利技术提出了一种基于强化学习内存调度决策的模型训练方法及系统，本发明专利技术基于强化学习算法，对于训练过程中产生的训练信息进行分析，并根据反馈更新决策方案，决定对哪些数据进行转移。从而进一步优化内存空间，并提高深度学习模型训练的整体性能。度学习模型训练的整体性能。度学习模型训练的整体性能。

【技术实现步骤摘要】
一种基于强化学习内存调度决策的模型训练方法及系统


[0001]本专利技术涉及计算机科学人工智能领域，尤其涉及一种针对基于强化学习的深度学习GPU显存管理的内存调度决策的模型训练方法与系统。

技术介绍

[0002]深度学习技术的革新大大推动了计算机视觉，自然语言处理，医学等领域的发展。为了进一步提高准确率，深度学习模型规模不断增大，结构更加复杂，这对于训练设备的存储空间提出了更高的需求。然而，以GPU为代表的深度学习加速器的存储空间有限，限制了研究人员开发探索规模更大的网络，制约了深度学习技术的迭代。目前常用的内存优化方案在训练过程中选择将GPU显存内暂时不需要的部分层数据转移至系统内存，以减轻GPU显存压力，并在后期需要访问该数据前将所需数据转移回GPU显存。该技术的难点是选择哪些数据在什么阶段进行交换。然而，GPU
‑
CPU之间的转移带宽有限，粗粒度的交换大量数据使得转移时间较长，引入无法忽略的通信延迟。现有方案通过使用传统启发式方法或专家经验进行转移策略设计。但是，难以在大规模的神经网络上决策出较优的数据转移方案，从而带来严重的模型训练性能损耗。

技术实现思路

[0003]本专利技术的目的是针对现有技术的不足，提出一种基于强化学习内存调度决策的模型训练方法及系统。
[0004]本专利技术采用的技术方案具体如下：
[0005]一种基于强化学习内存调度决策的模型训练方法，包括以下步骤：
[0006]步骤一：在第一轮迭代训练过程中，初始化调度方案；
[0007]步骤二：采集上一轮次神经网络模型迭代训练完后的GPU的内存空间、转移带宽，以及神经网络模型的网络层信息，记录数据的大小、各层的依赖关系；并根据采集的信息计算数据交换时间和训练时间；
[0008]步骤三：基于步骤二获取的信息动态决策每轮次的调度方案，包括以下子步骤：
[0009](3.1)对步骤二获得的训练时间、内存空间消耗进行归一化处理后，进行加权评估作为上一轮次的奖励收益。
[0010](3.2)根据步骤(3.1)获得的上一轮次的奖励收益和上一轮次获得的反馈值Q采用ε
‑
greedy算法，选择当前轮次等待交换的张量和当前轮次的调度方案，神经网络模型依据调度方案进行当前轮次训练；
[0011](3.3)根据上一轮次的反馈值Q、奖励收益以及当前轮次的调度方案，对反馈值Q进行更新并存储至Q表；
[0012]步骤四：重复步骤二、三进行迭代训练；随着迭代次数增加，调度方案将收敛到调度最优解；当达到预期训练效果时，结束训练。
[0013]进一步地，所述步骤(3.1)中，当前轮次的奖励收益R
plan
具体如下：
[0014]R
plan
＝R
time
+w(R
mem
‑
R
time
)
[0015]其中，R
time
表示上一轮次的训练时间，R
mem
表示上一轮次的训练消耗的GPU内存空间，w表示在奖励收益R
plan
中，上一轮次的训练消耗的GPU内存空间R
mem
所占据的比例。
[0016]进一步地，所述步骤(3.3)中，对反馈值Q进行更新具体为：
[0017]Q
k
(s,t)＝Q
k
‑1(s,t)+α[R
plan
+γQ
k
‑1(s',t')
‑
Q
k
‑1(s,t)][0018]其中，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表查找获得。
[0019]进一步地，所述步骤(3.2)中，采用ε
‑
greedy算法，选择当前轮次等待交换的张量和当前轮次的调度方案时，每轮次只选择同一种网络层类型对应的张量作为等待交换的张量。
[0020]进一步地，所述模型训练方法包括预训练和正式训练，其中，预训练采用部分训练数据进行训练并动态决策每轮次的调度方案，正式训练采用全部训练数据在预训练的基础上继续进行训练并动态决策每轮次的调度方案。
[0021]一种基于强化学习内存调度决策的模型训练系统，用于实现上述基于强化学习内存调度决策的模型训练方法，包括：
[0022]采集模块，用于采集上一轮次神经网络模型迭代训练完后的GPU的内存空间、转移带宽，以及神经网络模型的网络层信息，记录数据的大小、各层的依赖关系；并根据采集的信息计算数据交换的对应时间和训练时间；
[0023]决策模块，用于基于采集模块获取的信息动态决策每轮次的调度方案；
[0024]训练模块，依据采集模块和决策模块的数据进行迭代训练；随着迭代次数增加，调度方案将收敛到调度最优解；当达到预期训练效果时，结束训练。
[0025]本专利技术的有益效果是：本专利技术基于强化学习算法，对于训练过程中产生的训练信息进行分析，并根据反馈更新决策方案，决定对哪些数据进行转移。从而进一步优化内存空间，并提高深度学习模型训练的整体性能。
附图说明
[0026]图1是系统架构图
[0027]图2是决策过程流程图
[0028]图3是CIFAR10上训练支持的最大批量增长比例统计图
[0029]图4是子数据集批量学习优化效果结果图
具体实施方式
[0030]下面结合附图和具体实施例对本专利技术作进一步说明。
[0031]本专利技术提供一种基于强化学习内存调度决策的模型训练方法，如图2所示，包括以下步骤：
[0032]步骤一：用户通过接口完成模型代码的实现后，在第一轮迭代训练过程中，深度学
习框架会初始化调度方案s并进行训练。调度方案描述了每轮次迭代中，训练数据中每个张量是否在迭代中被交换，包括选择的被交换的张量和不交换的张量；一般地，调度方案可以采用二维数组表示，即将数据的状态使用二维数组进行表示，以便于策略拓展。在调度方案中，每个二维数组的维度为2，其中，调度策略为0表示保留在GPU内存中，为1表示进行数据交换。一系列数组构成了一个状态，决策的变化过程体现为状态间的转移。系统的最终目标，是求解得到一组数组，从而达到理想的训练效果。初始化调度方案s可以自行设定。
[0033]步本文档来自技高网...

【技术保护点】

【技术特征摘要】
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')

【专利技术属性】
技术研发人员：何水兵，宗威旭，陈平，汪睿，
申请(专利权)人：浙江大学，
类型：发明
国别省市：