CN112882830A - 显存管理、模型训练方法、装置、电子设备及存储介质 - Google Patents

Abstract

本发明公开一种显存管理、模型训练方法、装置、电子设备及存储介质，该显存管理方法包括：获取模型的当前轮训练对应的显存阈值；在所述电子设备的显存占用值大于所述显存阈值的情况下，确定满足张量选择规则的目标张量；释放所述目标张量占用的显存。可见，实施上述方法，可以通过释放显存中不影响模型正常训练的张量，来减少GPU显存占用，而不需要提前获得整张计算图的全局信息，做到了完全动态，从而实现动态图机制的深度学习框架下的显存管理。

Description

显存管理、模型训练方法、装置、电子设备及存储介质

技术领域

本发明涉及深度学习技术领域，特别涉及一种显存管理、模型训练方法、装置、电子设备及存储介质。

背景技术

在深度学习领域中，随着训练数据的增加，模型的尺寸和复杂度也大大增加，目前在模型训练过程中，经常会遇到有限的GPU(Graphics Processing Unit，图像处理器)显存资源无法满足Batch Size(批大小)较大的模型训练的困境，这就给深度学习框架带来了新的挑战，能否在模型训练时有效利用有限的计算存储资源，尤其是减少GPU显存占用，是评估深度学习框架性能的重要指标。

相关技术中，深度学习框架有以下几种降低显存占用的方法：通过合适的梯度定义，使得算子例如Relu、Sigmoid等的反向梯度计算不再依赖前向计算作为输入，从而在前向计算完成后就可以释放这部分显存；或者，计算每个算子的生命周期，生命周期不重叠的算子可以共享显存；或者，通过额外的数据传输减少显存占用，例如把暂时不用的数据从GPU交换到CPU(Central Processing Unit，中央处理器)，需要时再从CPU交换回来；或者，通过额外的计算减少显存占用，例如利用梯度检查点(Gradient Check pointing)重新计算中间结果的亚线性显存优化方法。

然而，相关技术中的以上方法均需事先获得计算图的全局信息，这就要求深度学习框架的计算图必须是静态图，而这对于动态图机制的深度学习框架是不可用的，因此，提出一种适用于动态图机制的深度学习框架的显存管理方法，已成为本领域技术人员亟待解决的技术问题。

发明内容

本发明实施例提供一种显存管理、模型训练方法、装置、电子设备及存储介质，以解决现有技术中的显存管理方法在动态图机制的深度学习框架上不可用的技术问题。

根据本发明的第一方面，公开了一种显存管理方法，所述方法包括：

获取模型的当前轮训练对应的显存阈值；

在所述电子设备的显存占用值大于所述显存阈值的情况下，确定满足张量选择规则的目标张量；

释放所述目标张量占用的显存。

可选地，作为一些实施例，所述确定满足张量选择规则的目标张量，包括：

根据目标估值函数，计算所述显存中的张量对应的估值函数值；

将估值函数值最大的张量确定为所述目标张量。

可选地，作为一些实施例，所述根据目标估值函数，计算所述显存中的张量对应的估值函数值，包括：

根据目标估值函数，计算所述显存中未上锁的张量对应的估值函数值；

所述将估值函数值最大的张量确定为所述目标张量，包括：

将未上锁的且估值函数值最大的张量，确定为所述目标张量。

可选地，作为一些实施例，所述方法还包括：

根据张量占用的显存大小、张量占用显存的时长、张量的计算代价和张量的重计算次数中的至少两个或多个，确定所述目标估值函数。

可选地，作为一些实施例，所述目标估值函数为：

其中，t为张量，M(t)为t占用的显存大小，L(t)为t占用显存的时长，C(t)为t的计算代价，R(t)为t的重计算次数，α、β、γ和δ为所述目标估值函数的超参数。

可选地，作为一些实施例，在所述模型的当前轮训练为第一轮训练时，所述显存阈值＝显存容量/2，α＝β＝γ＝1，δ＝1/2。

可选地，作为一些实施例，所述方法还包括：

获取第N轮训练的显存占用值超过显存阈值的次数；

根据所述显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值，其中，N为大于等于1的整数。

可选地，作为一些实施例，所述根据所述显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值，包括：

在所述显存占用值超过显存阈值的次数大于第一次数阈值的情况下，增大第N+1轮训练对应的显存阈值；

在所述显存占用值超过显存阈值的次数不大于所述第一次数阈值的情况下，减小第N+1轮训练对应的显存阈值。

可选地，作为一些实施例，所述方法还包括：

获取第N轮训练的申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比；

根据所述申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值。

可选地，作为一些实施例，所述根据所述申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值，包括以下步骤中的至少一个：

在所述申请空间大于显存碎片的最大值的次数大于第二次数阈值的情况下，增大第N+1轮训练对应的α的取值；

在所述重新计算的时长占总训练时长的百分比相较于上一轮模型训练增加的情况下，减少第N+1轮训练对应的γ的取值。

可选地，作为一些实施例，所述方法还包括：

计算第N轮训练的耗时参数＝重新计算耗时/原本计算耗时；

基于所述耗时参数运行模拟退火算法，调整第N+1轮训练对应的显存阈值和/或所述目标估值函数中超参数的取值。

可选地，作为一些实施例，所述张量的计算代价的获取步骤包括：

从缓存中读取与所述张量具有相同算子和/或相同输入形状的历史张量的历史计算代价；

将读取到的所述历史计算代价确定为所述张量对应的计算代价。

可选地，作为一些实施例，所述张量占用显存的时长的获取步骤包括：

获取当前轮训练中正在执行的算子个数和所述张量进入显存的时间；

根据正在执行的算子个数和所述张量进入显存的时间，确定所述张量占用显存的时长。

可选地，作为一些实施例，所述方法还包括：

确定当前轮训练是否包含目标算子执行序列，若是，则将所述目标算子执行序列对应的张量确定为所述目标张量；

其中，所述目标算子执行序列中包括：按照特定顺序执行的多个算子，所述目标算子执行序列对应的张量在历史训练过程中有被释放的记录。

可选地，作为一些实施例，在所述释放所述目标张量占用的显存的步骤之后，还包括：

将所述当前轮训练中新产生的张量存储至所释放的显存中。

根据本发明的第二方面，公开了一种模型训练方法，所述方法包括：

获取训练样本集，其中，所述训练样本集中包括用于模型训练的训练数据；

基于所述训练样本集和初始模型进行模型训练，在模型训练过程中，基于第一方面中显存管理方法对训练产生的张量进行管理，直至训练得到目标模型。

根据本发明的第三方面，公开了一种显存管理装置，所述装置包括：

第一获取模块，用于获取模型的当前轮训练对应的显存阈值；

第一确定模块，用于在所述电子设备的显存占用值大于所述显存阈值的情况下，确定满足张量选择规则的目标张量；

释放模块，用于释放所述目标张量占用的显存。

可选地，作为一些实施例，所述第一确定模块包括：

计算子模块，用于根据目标估值函数，计算所述显存中的张量对应的估值函数值；

确定子模块，用于将估值函数值最大的张量确定为所述目标张量。

可选地，作为一些实施例，所述计算子模块包括：

计算单元，用于根据目标估值函数，计算所述显存中未上锁的张量对应的估值函数值；

所述确定子模块包括：

确定单元，用于将未上锁的且估值函数值最大的张量，确定为所述目标张量。

可选地，作为一些实施例，所述装置还包括：

第二确定模块，用于根据张量占用的显存大小、张量占用显存的时长、张量的计算代价和张量的重计算次数中的至少两个或多个，确定所述目标估值函数。

可选地，作为一些实施例，所述目标估值函数为：

其中，t为张量，M(t)为t占用的显存大小，L(t)为t占用显存的时长，C(t)为t的计算代价，R(t)为t的重计算次数，α、β、γ和δ为所述目标估值函数的超参数。

可选地，作为一些实施例，在所述模型的当前轮训练为第一轮训练时，所述显存阈值＝显存容量/2，α＝β＝γ＝1，δ＝1/2。

可选地，作为一些实施例，所述装置还包括：

第二获取模块，用于获取第N轮训练的显存占用值超过显存阈值的次数；

第一调整模块，用于根据所述显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值，其中，N为大于等于1的整数。

可选地，作为一些实施例，所述第一调整模块包括：

第一调整子模块，用于在所述显存占用值超过显存阈值的次数大于第一次数阈值的情况下，增大第N+1轮训练对应的显存阈值；

第二调整子模块，用于在所述显存占用值超过显存阈值的次数不大于所述第一次数阈值的情况下，减小第N+1轮训练对应的显存阈值。

可选地，作为一些实施例，所述装置还包括：

第三获取模块，用于获取第N轮训练的申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比；

第二调整模块，用于根据所述申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值。

可选地，作为一些实施例，所述第二调整模块包括以下子模块中的至少一个：

第三调整子模块，用于在所述申请空间大于显存碎片的最大值的次数大于第二次数阈值的情况下，增大第N+1轮训练对应的α的取值；

第四调整子模块，用于在所述重新计算的时长占总训练时长的百分比相较于上一轮模型训练增加的情况下，减少第N+1轮训练对应的γ的取值。

可选地，作为一些实施例，所述装置还包括：

计算模块，用于计算第N轮训练的耗时参数＝重新计算耗时/原本计算耗时；

第三调整模块，用于基于所述耗时参数运行模拟退火算法，调整第N+1轮训练对应的显存阈值和/或所述目标估值函数中超参数的取值。

可选地，作为一些实施例，所述装置还包括：

读取模块，用于从缓存中读取与所述张量具有相同算子和/或相同输入形状的历史张量的历史计算代价；

第三确定模块，用于将读取到的所述历史计算代价确定为所述张量对应的计算代价。

可选地，作为一些实施例，所述装置还包括：

第四获取模块，用于获取当前轮训练中正在执行的算子个数和所述张量进入显存的时间；

第四确定模块，用于根据正在执行的算子个数和所述张量进入显存的时间，确定所述张量占用显存的时长。

可选地，作为一些实施例，所述装置还包括：

第五确定模块，用于确定当前轮训练是否包含目标算子执行序列，若是，则将所述目标算子执行序列对应的张量确定为所述目标张量；

其中，所述目标算子执行序列中包括：按照特定顺序执行的多个算子，所述目标算子执行序列对应的张量在历史训练过程中有被释放的记录。

可选地，作为一些实施例，所述装置还包括：

存储模块，用于将所述当前轮训练中新产生的张量存储至所释放的显存中。

根据本发明的第四方面，公开了一种模型训练装置，所述装置包括：

第二获取模块，用于获取训练样本集，其中，所述训练样本集中包括用于模型训练的训练数据；

训练模块，用于基于所述训练样本集和初始模型进行模型训练，在模型训练过程中，基于第三方面中显存管理装置对训练产生的张量进行管理，直至训练得到目标模型。

根据本发明的第五方面，公开了一种电子设备，包括：存储器、处理器及存储在所述存储器上并可在所述处理器上运行的程序，所述程序被所述处理器执行时实现如第一方面中显存管理方法的步骤或者实现如第二方面中模型训练方法的步骤。

根据本发明的第六方面，公开了一种计算机可读存储介质，所述计算机可读存储介质上存储有程序，所述程序被所述处理器执行时实现如第一方面中显存管理方法的步骤或者实现如第二方面中模型训练方法的步骤。

本发明实施例中，可以获取模型的当前轮训练对应的显存阈值，在电子设备的显存占用值大于显存阈值的情况下，确定满足张量选择规则的目标张量，释放目标张量占用的显存。可见，本发明实施例中，可以通过释放显存中不影响模型正常训练的张量，来减少GPU显存占用，而不需要提前获得整张计算图的全局信息，做到了完全动态，从而实现动态图机制的深度学习框架下的显存管理。此外，在该动态图机制的深度学习框架下，用户不需要修改原来训练的代码，就可以显存容量不变的机器上以原来的耗时的常数倍的时间训练比原来大一倍的模型，模型训练效率较高。

附图说明

图1是本发明的一些实施例的显存管理方法的流程图；

图2是本发明的一些实施例的Python中的张量与C++中的张量的对应关系的示例图；

图3是本发明的一些实施例的当显存中只能存储三个张量时显存中张量变化情况的示例图；

图4是本发明的一些实施例的模型训练方法的流程图；

图5是本发明的一些实施例的显存管理装置的结构示意图；

图6是本发明的一些实施例的模型训练装置的结构示意图。

具体实施方式

为使本发明的上述目的、特征和优点能够更加明显易懂，下面结合附图和具体实施方式对本发明作进一步详细的说明。

需要说明的是，对于方法实施例，为了简单描述，故将其都表述为一系列的动作组合，但是本领域技术人员应该知悉，本发明实施例并不受所描述的动作顺序的限制，因为依据本发明实施例，某些步骤可以采用其他顺序或者同时进行。其次，本领域技术人员也应该知悉，说明书中所描述的实施例均属于优选实施例，所涉及的动作并不一定是本发明实施例所必须的。

在深度学习领域中，随着训练数据的增加，模型的尺寸和复杂度也大大增加，我们经常遇到的困境是有限的GPU显存资源无法满足Batch Size较大的模型训练。这就给深度学习框架带来了新的挑战，能否在模型训练时有效利用有限的计算存储资源，尤其是减少GPU显存占用，是评估深度学习框架性能的重要指标。

相关技术中，在计算存储资源一定的情况下，深度学习框架有几种降低显存占用的方法，具体如下：

方法一、通过合适的梯度定义，使得算子例如Relu、Sigmoid的反向梯度计算不再依赖前向计算作为输入，从而在前向计算完成后就可以释放这部分内存；

方法二、计算每个算子的生命周期，生命周期不重叠的算子可以共享显存；

方法三、通过额外的数据传输减少显存占用，例如，把暂时不用的数据从GPU交换到CPU，需要时再从CPU交换回来；

方法四、通过额外的计算减少显存占用，例如，利用梯度检查点(Gradient Checkpointing)重新计算中间结果的亚线性显存优化方法。

以上几种方法的缺陷在于，需要事先获得计算图的全局信息，这就要求计算图必须是静态的。静态图的好处是程序在编译执行时就能生成神经网络的结构，这种机制允许编译器进行最大程度的优化，例如，计算每个算子的生命周期从而更好地分配显存，通过搜索算法计算最优的梯度检查点等等。但是，这也意味着编译器实际执行与用户期望的程序执行存在着很大的代沟，代码中的错误将更加难以发现。与之相反，动态图下的程序会完全按照用户编写命令的顺序执行，易于调试，对用户更加友好。因此，考虑到便捷性，用户更多地会去选择使用具有动态图机制的灵活框架，例如，PyTorch、MegEngine等等。

也正是因为动态图灵活的特性，使得上述的技术不能很好地从静态图迁移到动态图上。例如，方法二在动态图上不可行，因为动态图上无法获得未来的计算序列，也就不可能计算每个算子的生命周期；方法四在动态图上不可行，因为动态图不能在一开始就获得整张计算图的信息，所以不能计算出最优的梯度检查点。可见，以上的显存管理方法，对于动态图机制的深度学习框架不可用。

为解决上述技术问题，本发明实施例提供了一种显存管理、模型训练方法、装置、电子设备及存储介质。

下面首先对本发明实施例提供的一种显存管理方法进行介绍。

需要说明的是，本发明实施例提供的显存管理方法适应于电子设备，在实际应用中，该电子设备可以为服务器等，本发明实施例对此不作限定。

图1是本发明的一些实施例的显存管理方法的流程图，如图1所示，该方法可以包括以下步骤：步骤101、步骤102和步骤103，其中，

在步骤101中，获取模型的当前轮训练对应的显存阈值。

为便于理解，结合一个例子对本发明实施例中涉及到的深度学习框架和基于该深度学习框架进行模型训练的相关内容进行描述。

在进行实际的模型训练之前，用户会在执行所述模型训练的机器(主要是服务器，后续均以服务器为例进行描述)上部署支持动态图的深度学习框架，以及会用Python语言编写相应的模型代码，例如，如果用户期望最终训练得到一个对象分类模型，那么需要编写对象分类模型对应的Python代码。之后，用户将编写的模型代码提交到服务器的深度学习框架上运行，深度学习框架会提供一个解释器，该解释器会将模型代码拆分成多个小任务，发送给服务器的GPU进行模型训练，例如，将模型代码中的核函数发射给GPU。在模型训练过程中，模型训练相关的一些数据会占用GPU的显存，例如，模型训练过程中产生的张量(Tensor)会存储在显存中，本发明实施例的目的在于：优化模型训练过程中显存占用情况。

本发明实施例中，解释器支持的一个功能就是将模型代码中的张量一一绑定在C++代码中的张量类型上，当需要将GPU显存上的某一张量占用的显存释放时，只需要在C++中操作张量结构体就可以实现。

如图2所示，图2示出了Python代码中的张量与C++中的张量的对应关系，以张量为t为例，可以在C++的张量类型中记录一些辅助信息：计算历史信息和属性值，其中，计算历史信息包括：计算出t的算子以及这个算子的所有输入张量，属性值包括：M(t)为t占用的显存大小(单位：MB)、L(t)为t占用显存的时长、C(t)为t的计算代价(单位：ms)，以及R(t)为t的重计算次数。

考虑到模型训练过程中，基础模型(又称为“初始模型”或“骨干网络”等)需要经过多轮训练才能得到目标模型，而每轮模型训练过程中对显存的占用情况通常不同，本发明实施例中，可以以模型训练的每一轮为处理单位，针对每一轮模型训练，可以根据模型和显存的占用情况等，采用相同/不同的显存优化参数，对每一轮训练过程中的显存占用进行优化。

本发明实施例中，在模型训练过程中，每开始一轮训练，会获取模型的当前训练轮数，也就是，模型当前训练到第几轮。

为了防止显存耗尽，本发明实施例中，针对每一轮训练，设定一个显存阈值，当电子设备的显存占用值超过这个显存阈值时，根据张量选择规则，从显存中存储的张量中查找不会影响模型后续训练过程的目标张量，将它从显存中释放掉，重复上述查找和释放操作，直到电子设备的显存占用值低于显存阈值。

本发明实施例中，对于释放的张量，如果用户未来需要访问这个张量，或者在反向求导时需要用到这个张量，则可以先根据该张量的计算路径(存储在CPU内存中)、算子和所有输入，恢复出原来的值。也就是，只需要在框架中记录每个张量的计算路径，寻找张量、释放显存、重计算恢复张量等细节是不需要用户来实现的，用户对整个优化的过程是无感知的，不会影响到用户的编程体验。

本发明实施例中，在模型训练过程中，为了保证显存优化效果，显存阈值和张量选择规则，会随着模型的训练做一定的动态调整，也就是，针对一个模型的训练过程，不同的训练轮数，其对应的显存阈值和张量选择规则通常是不同的。因此，在本轮模型训练过程中，在优化显存占用时，需要获取当前轮训练对应的显存阈值和张量选择规则。

在步骤102中，在电子设备的显存占用值大于显存阈值的情况下，确定满足张量选择规则的目标张量。

在本发明提供的一些实施方式中，张量选择规则可以基于张量的估值函数实现，此时，上述步骤102具体可以包括以下步骤(图中未示出)：步骤1021和步骤1022，其中，

在步骤1021中，在电子设备的显存占用值大于显存阈值的情况下，根据目标估值函数，计算显存中的张量对应的估值函数值；

在步骤1022中，将估值函数值最大的张量确定为目标张量。

本发明实施例中，可以根据张量占用的显存大小、张量占用显存的时长、张量的计算代价和张量的重计算次数中的至少两个或多个，确定目标估值函数。

本发明实施例中，当根据张量占用的显存大小、张量占用显存的时长、张量的计算代价和张量的重计算次数确定目标估值函数时，目标估值函数可以为：

其中，t为张量，M(t)为t占用的显存大小，L(t)为t占用显存的时长，C(t)为t的计算代价，R(t)为t的重计算次数，α、β、γ和δ为目标估值函数的超参数。

本发明实施例中，在模型的当前轮训练为第一轮训练时，显存阈值＝显存容量/2，α＝β＝γ＝1，δ＝1/2。也就是，在模型开始第一轮训练时，初始策略为：设定显存阈值＝显存容量/2，设置超参数α＝β＝γ＝1，δ＝1/2。当然本发明实施对在模型开始第一轮训练时的初始策略不作限定，还可以根据实际应用需求，相应设定其他初始策略。

本发明实施例中，可以根据模型的历史训练信息，对每一轮训练对应的显存阈值和目标估值函数的超参数进行动态微调。

基于上述动态微调思路，本发明实施例提供的显存管理方法，还可以包括以下步骤：

获取第N轮训练的显存占用值超过显存阈值的次数；

根据显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值，其中，N为大于等于1的整数。

本发明实施例中，在根据显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值时，可以采用以下策略：

在显存占用值超过显存阈值的次数大于第一次数阈值的情况下，增大第N+1轮训练对应的显存阈值；

在显存占用值超过显存阈值的次数不大于第一次数阈值的情况下，减小第N+1轮训练对应的显存阈值。

在一个例子中，在显存占用值超过显存阈值的次数大于第一次数阈值的情况下，设置：第N+1轮训练对应的显存阈值＝第N轮训练对应的显存阈值*(1+10％)；在显存占用值超过显存阈值的次数不大于第一次数阈值的情况下，设置：第N+1轮训练对应的显存阈值＝第N轮训练对应的显存阈值*(1-10％)。

基于上述动态微调思路，本发明实施例提供的显存管理方法，还可以包括以下步骤：

获取第N轮训练的申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比；

根据申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值。

本发明实施例中，在根据申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值时，可以采用以下策略：

在申请空间大于显存碎片的最大值的次数大于第二次数阈值的情况下，增大第N+1轮训练对应的α的取值；

在重新计算的时长占总训练时长的百分比相较于上一轮模型训练增加的情况下，减少第N+1轮训练对应的γ的取值。

在一个例子中，在申请空间大于显存碎片的最大值的次数大于第二次数阈值的情况下，说明应当更侧重于释放显存占用较大的张量，此时，设置：第N+1轮训练对应的α＝第N轮训练对应的α*(1+10％)；在重新计算的时长占总训练时长的百分比相较于上一轮模型训练增加的情况下，说明应当更侧重于释放重计算代价较小的张量，此时，设置：第N+1轮训练对应的γ＝第N轮训练对应的γ*(1-10％)。

除了前述调整策略，考虑到每一轮的训练耗时主要由原本计算耗时和重新计算耗时两部分组成，为了在不超过显存限制的前提下成功训练该模型，引入耗时参数P＝重新计算耗时/原本计算耗时，并使得P尽可能小。在一些实施例中，可以基于模拟退火算法对超参数进行调整，本发明实施例对此不作限定。此时，本发明实施例提供的显存管理方法，还可以增加以下步骤：

计算第N轮训练的耗时参数＝重新计算耗时/原本计算耗时；

基于耗时参数运行模拟退火算法，调整第N+1轮训练对应的显存阈值和/或目标估值函数中超参数的取值。

由于模拟退火算法为现有算法，在此不再赘述。

本发明实施例中，考虑到在测量张量的计算代价C(t)时，由于会引入设备间的同步，应当尽量避免测量次数过多，造成额外的过大的开销。针对这种情况，可以引入Cache(缓存)优化，对于相同算子和/或相同输入形状的张量，只测量一次耗时，并存入Cache中，之后遇到相同算子和相同输入形状的张量，直接从Cache中读取耗时即可。

相应的，基于上述思路，张量的计算代价的获取步骤，可以包括以下步骤：从缓存中读取与张量具有相同算子和/或相同输入形状的历史张量的历史计算代价；将读取到的历史计算代价确定为该张量对应的计算代价。

本发明实施例中，在计算张量占用显存的时长L(t)时，为了避免大量的调用获取当前时间的函数，改用当前执行的算子个数作为时间戳来近似计算L(t)＝当前时间戳-张量进入显存时的时间戳，以减少调用时间函数的开销。

相应的，基于上述思路，张量占用显存的时长的获取步骤，可以包括以下步骤：获取当前轮训练中正在执行的算子个数和张量进入显存的时间；根据正在执行的算子个数和张量进入显存的时间，确定该张量占用显存的时长。

本发明实施例中，在发射一个算子的核函数前，需要对这个算子依赖的输入张量上锁，防止其被确定为目标张量，那将导致设备在计算该算子时，显存中不存在这个张量，在算子计算完成后，可以把输入张量的锁释放掉。

针对上述情况，上述步骤1021具体可以包括以下步骤：

根据目标估值函数，计算显存中未上锁的张量对应的估值函数值；

相应的，上述步骤1022具体可以包括以下步骤：

将未上锁的且估值函数值最大的张量，确定为目标张量。

本发明实施例中，当用户在模型代码中删除张量t时，解释器需要判断依赖t作为算子输入的张量是否都被删除了，只要有一个张量u依赖t，就不能在解释器中删除t，因为将来u可能作为目标张量，需要t的信息进行恢复，这时只需要在解释器中释放掉t的显存，不需要真正删除。也就是，只有当依赖t的张量全部被用户删除时，解释器才能删除t。

在步骤103中，释放目标张量占用的显存。

本发明实施例中，在将目标张量占用的显存释放后，可以将当前轮训练中新产生的张量存储至所释放的显存中，以保证模型训练能够正常进行。

在一个例子中，如图3所示，图3示出了显存容量为三个张量时，显存中的张量的变化情况。在模型训练过程中，显存中存储有A、B和C，三个张量，此时，电子设备的显存占用值为100％，如果显存阈值为70％，由于显存占用值100％大于显存阈值70％，因此需要查找出显存中的目标张量，例如若C为目标张量，则将C占用的显存释放，以达到优化显存的目的。

由上述实施例可见，该实施例中，可以获取模型的当前轮训练对应的显存阈值，在电子设备的显存占用值大于显存阈值的情况下，确定满足张量选择规则的目标张量，释放目标张量占用的显存。可见，本发明实施例中，可以通过释放显存中不影响模型正常训练的张量，来减少GPU显存占用，而不需要提前获得整张计算图的全局信息，做到了完全动态，从而实现动态图机制的深度学习框架下的显存管理。此外，在该动态图机制的深度学习框架下，用户不需要修改原来训练的代码，就可以显存容量不变的机器上以原来的耗时的常数倍的时间训练比原来大一倍的模型，例如：在一张2080Ti(显存11GB)卡上，ResNet50的训练批大小可以变成原来的两倍，达到250左右，模型训练效率较高。

在本发明提供的一些实施例中，显存管理方法还可以在图1所示实施例的基础上，增加以下步骤：确定当前轮训练是否包含目标算子执行序列，若是，则将目标算子执行序列对应的张量确定为目标张量；

其中，目标算子执行序列中包括：按照特定顺序执行的多个算子，目标算子执行序列对应的张量在历史训练过程中有被释放的记录。

在一个例子中，如果形如“卷积-批正则化-Relu-池化-卷积-批正则化-Relu”，长度为7的计算序列对应的张量都在之后的某个时刻全部被释放，那么当在后续训练过程中再次发现刚刚执行了这7个算子的时候，可以主动将其对应的张量所占用的显存释放，这样就可以避免在将来显存超过阈值的时候再被动地进行查找，降低搜索的开销。

图4是本发明的一些实施例的模型训练方法的流程图，如图4所示，该方法可以包括以下步骤：步骤401和步骤402，其中，

在步骤401中，获取训练样本集，其中，训练样本集中包括用于模型训练的训练数据。

在步骤402中，基于训练样本集和初始模型进行模型训练，在模型训练过程中，基于预设显存管理方法对训练产生的张量进行管理，直至训练得到目标模型。

本发明实施例中，预设显存管理方法为上述任一一些实施例中的显存管理方法。

本发明实施例中，目标模型包括但不限于以下用途的模型：用于确定待处理图像所属类别的模型、用于识别待处理图像中人脸的模型、用于检测待处理图像中特定对象的模型、用于分割待处理图像中特定对象的模型，以及用于生成新图像的模型等等，其中，新图像具有与待处理图像相似的特征。

本申请实施例中，可以根据目标模型的用途，选择合适的训练样本集。

由上述实施例可见，该实施例中，可以基于上述显存管理方法对模型训练过程中产生的张量进行管理，使得用户不需要修改原来训练的代码，就可以显存容量不变的机器上以原来的耗时的常数倍的时间训练比原来大一倍的模型，例如：在一张2080Ti(显存11GB)卡上，ResNet50的训练批大小可以变成原来的两倍，达到250左右，模型训练效率较高。

图5是本发明的一些实施例的显存管理装置的结构示意图，如图5所示，显存管理装置500可以包括：第一获取模块501、第一确定模块502和释放模块503，其中，

第一获取模块501，用于获取模型的当前轮训练对应的显存阈值；

第一确定模块502，用于在所述电子设备的显存占用值大于所述显存阈值的情况下，确定满足张量选择规则的目标张量；

释放模块503，用于释放所述目标张量占用的显存。

由上述实施例可见，该实施例中，可以获取模型的当前轮训练对应的显存阈值，在电子设备的显存占用值大于显存阈值的情况下，确定满足张量选择规则的目标张量，释放目标张量占用的显存。可见，本发明实施例中，可以通过释放显存中不影响模型正常训练的张量，来减少GPU显存占用，而不需要提前获得整张计算图的全局信息，做到了完全动态，从而实现动态图机制的深度学习框架下的显存管理。此外，在该动态图机制的深度学习框架下，用户不需要修改原来训练的代码，就可以显存容量不变的机器上以原来的耗时的常数倍的时间训练比原来大一倍的模型，模型训练效率较高。

可选地，作为一些实施例，所述第一确定模块502，可以包括：

计算子模块，用于根据目标估值函数，计算所述显存中的张量对应的估值函数值；

确定子模块，用于将估值函数值最大的张量确定为所述目标张量。

可选地，作为一些实施例，所述计算子模块，可以包括：

计算单元，用于根据目标估值函数，计算所述显存中未上锁的张量对应的估值函数值；

所述确定子模块，可以包括：

确定单元，用于将未上锁的且估值函数值最大的张量，确定为所述目标张量。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

第二确定模块，用于根据张量占用的显存大小、张量占用显存的时长、张量的计算代价和张量的重计算次数中的至少两个或多个，确定所述目标估值函数。

可选地，作为一些实施例，所述目标估值函数为：

其中，t为张量，M(t)为t占用的显存大小，L(t)为t占用显存的时长，C(t)为t的计算代价，R(t)为t的重计算次数，α、β、γ和δ为所述目标估值函数的超参数。

可选地，作为一些实施例，在所述模型的当前轮训练为第一轮训练时，所述显存阈值＝显存容量/2，α＝β＝γ＝1，δ＝1/2。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

第二获取模块，用于获取第N轮训练的显存占用值超过显存阈值的次数；

第一调整模块，用于根据所述显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值，其中，N为大于等于1的整数。

可选地，作为一些实施例，所述第一调整模块，可以包括：

第一调整子模块，用于在所述显存占用值超过显存阈值的次数大于第一次数阈值的情况下，增大第N+1轮训练对应的显存阈值；

第二调整子模块，用于在所述显存占用值超过显存阈值的次数不大于所述第一次数阈值的情况下，减小第N+1轮训练对应的显存阈值。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

第三获取模块，用于获取第N轮训练的申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比；

第二调整模块，用于根据所述申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值。

可选地，作为一些实施例，所述第二调整模块，可以包括以下子模块中的至少一个：

第三调整子模块，用于在所述申请空间大于显存碎片的最大值的次数大于第二次数阈值的情况下，增大第N+1轮训练对应的α的取值；

第四调整子模块，用于在所述重新计算的时长占总训练时长的百分比相较于上一轮模型训练增加的情况下，减少第N+1轮训练对应的γ的取值。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

计算模块，用于计算第N轮训练的耗时参数＝重新计算耗时/原本计算耗时；

第三调整模块，用于基于所述耗时参数运行模拟退火算法，调整第N+1轮训练对应的显存阈值和/或所述目标估值函数中超参数的取值。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

读取模块，用于从缓存中读取与所述张量具有相同算子和/或相同输入形状的历史张量的历史计算代价；

第三确定模块，用于将读取到的所述历史计算代价确定为所述张量对应的计算代价。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

第四获取模块，用于获取当前轮训练中正在执行的算子个数和所述张量进入显存的时间；

第四确定模块，用于根据正在执行的算子个数和所述张量进入显存的时间，确定所述张量占用显存的时长。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

第五确定模块，用于确定当前轮训练是否包含目标算子执行序列，若是，则将所述目标算子执行序列对应的张量确定为所述目标张量；

其中，所述目标算子执行序列中包括：按照特定顺序执行的多个算子，所述目标算子执行序列对应的张量在历史训练过程中有被释放的记录。

可选地，作为一些实施例，所述显存管理装置500，还可以包括：

存储模块，用于将所述当前轮训练中新产生的张量存储至所释放的显存中。

图6是本发明的一些实施例的模型训练装置的结构示意图，如图6所示，模型训练装置600可以包括：第二获取模块601和训练模块602，其中，

第二获取模块601，用于获取训练样本集，其中，所述训练样本集中包括用于模型训练的训练数据；

训练模块602，用于基于所述训练样本集和初始模型进行模型训练，在模型训练过程中，基于上述任一显存管理装置对训练产生的张量进行管理，直至训练得到目标模型。

由上述实施例可见，该实施例中，可以基于上述显存管理方法对模型训练过程中产生的张量进行管理，使得用户不需要修改原来训练的代码，就可以显存容量不变的机器上以原来的耗时的常数倍的时间训练比原来大一倍的模型，例如：在一张2080Ti(显存11GB)卡上，ResNet50的训练批大小可以变成原来的两倍，达到250左右，模型训练效率较高。

对于装置实施例而言，由于其与方法实施例基本相似，所以描述的比较简单，相关之处参见方法实施例的部分说明即可。

根据本发明的一些实施例，本发明还提供了一种电子设备，包括：存储器、处理器及存储在所述存储器上并可在所述处理器上运行的程序，所述程序被所述处理器执行时实现如上述任意一些实施例所述的显存管理方法中的步骤或者实现如上述一些实施例所述的模型训练方法中的步骤。

根据本发明的一些实施例，本发明还提供了一种计算机可读存储介质，所述计算机可读存储介质上存储有程序，所述程序被处理器执行时实现如上述任意一些实施例所述的显存管理方法中的步骤或者实现如上述一些实施例所述的模型训练方法中的步骤。

本说明书中的各个实施例均采用递进的方式描述，每个实施例重点说明的都是与其他实施例的不同之处，各个实施例之间相同相似的部分互相参见即可。

本领域内的技术人员应明白，本发明实施例的实施例可提供为方法、装置、或计算机程序产品。因此，本发明实施例可采用完全硬件实施例、完全软件实施例、或结合软件和硬件方面的实施例的形式。而且，本发明实施例可采用在一个或多个其中包含有计算机可用程序代码的计算机可用存储介质(包括但不限于磁盘存储器、CD-ROM、光学存储器等)上实施的计算机程序产品的形式。

本发明实施例是参照根据本发明实施例的方法、终端设备(***)、和计算机程序产品的流程图和/或方框图来描述的。应理解可由计算机程序指令实现流程图和/或方框图中的每一流程和/或方框、以及流程图和/或方框图中的流程和/或方框的结合。可提供这些计算机程序指令到通用计算机、专用计算机、嵌入式处理机或其他可编程数据处理终端设备的处理器以产生一个机器，使得通过计算机或其他可编程数据处理终端设备的处理器执行的指令产生用于实现在流程图一个流程或多个流程和/或方框图一个方框或多个方框中指定的功能的装置。

这些计算机程序指令也可存储在能引导计算机或其他可编程数据处理终端设备以特定方式工作的计算机可读存储器中，使得存储在该计算机可读存储器中的指令产生包括指令装置的制造品，该指令装置实现在流程图一个流程或多个流程和/或方框图一个方框或多个方框中指定的功能。

这些计算机程序指令也可装载到计算机或其他可编程数据处理终端设备上，使得在计算机或其他可编程终端设备上执行一系列操作步骤以产生计算机实现的处理，从而在计算机或其他可编程终端设备上执行的指令提供用于实现在流程图一个流程或多个流程和/或方框图一个方框或多个方框中指定的功能的步骤。

尽管已描述了本发明实施例的优选实施例，但本领域内的技术人员一旦得知了基本创造性概念，则可对这些实施例做出另外的变更和修改。所以，所附权利要求意欲解释为包括优选实施例以及落入本发明实施例范围的所有变更和修改。

最后，还需要说明的是，在本文中，诸如第一和第二等之类的关系术语仅仅用来将一个实体或者操作与另一个实体或操作区分开来，而不一定要求或者暗示这些实体或操作之间存在任何这种实际的关系或者顺序。而且，术语“包括”、“包含”或者其任何其他变体意在涵盖非排他性的包含，从而使得包括一系列要素的过程、方法、物品或者终端设备不仅包括那些要素，而且还包括没有明确列出的其他要素，或者是还包括为这种过程、方法、物品或者终端设备所固有的要素。在没有更多限制的情况下，由语句“包括一个……”限定的要素，并不排除在包括所述要素的过程、方法、物品或者终端设备中还存在另外的相同要素。

以上对本发明所提供的一种显存管理、模型训练方法、装置、电子设备及存储介质，进行了详细介绍，本文中应用了具体个例对本发明的原理及实施方式进行了阐述，以上实施例的说明只是用于帮助理解本发明的方法及其核心思想；同时，对于本领域的一般技术人员，依据本发明的思想，在具体实施方式及应用范围上均会有改变之处，综上所述，本说明书内容不应理解为对本发明的限制。

Claims (20)

1.一种显存管理方法，应用于电子设备，其特征在于，所述方法包括：

获取模型的当前轮训练对应的显存阈值；

在所述电子设备的显存占用值大于所述显存阈值的情况下，确定满足张量选择规则的目标张量；

释放所述目标张量占用的显存。

2.根据权利要求1所述的方法，其特征在于，所述确定满足张量选择规则的目标张量，包括：

根据目标估值函数，计算所述显存中的张量对应的估值函数值；

将估值函数值最大的张量确定为所述目标张量。

3.根据权利要求2所述的方法，其特征在于，所述根据目标估值函数，计算所述显存中的张量对应的估值函数值，包括：

根据目标估值函数，计算所述显存中未上锁的张量对应的估值函数值；

所述将估值函数值最大的张量确定为所述目标张量，包括：

将未上锁的且估值函数值最大的张量，确定为所述目标张量。

4.根据权利要求2或3所述的方法，其特征在于，所述方法还包括：

根据张量占用的显存大小、张量占用显存的时长、张量的计算代价和张量的重计算次数中的至少两个或多个，确定所述目标估值函数。

5.根据权利要求4所述的方法，其特征在于，所述目标估值函数为：

其中，t为张量，M(t)为t占用的显存大小，L(t)为t占用显存的时长，C(t)为t的计算代价，R(t)为t的重计算次数，α、β、γ和δ为所述目标估值函数的超参数。

6.根据权利要求5所述的方法，其特征在于，在所述模型的当前轮训练为第一轮训练时，所述显存阈值＝显存容量/2，α＝β＝γ＝1，δ＝1/2。

7.根据权利要求1-6任一项所述的方法，其特征在于，所述方法还包括：

获取第N轮训练的显存占用值超过显存阈值的次数；

根据所述显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值，其中，N为大于等于1的整数。

8.根据权利要求7所述的方法，其特征在于，所述根据所述显存占用值超过显存阈值的次数，调整第N+1轮训练对应的显存阈值，包括：

在所述显存占用值超过显存阈值的次数大于第一次数阈值的情况下，增大第N+1轮训练对应的显存阈值；

在所述显存占用值超过显存阈值的次数不大于所述第一次数阈值的情况下，减小第N+1轮训练对应的显存阈值。

9.根据权利要求5-8任一项所述的方法，其特征在于，所述方法还包括：

获取第N轮训练的申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比；

根据所述申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值。

10.根据权利要求9所述的方法，其特征在于，所述根据所述申请空间大于显存碎片的最大值的次数和/或重新计算的时长占总训练时长的百分比，调整第N+1轮训练对应的目标估值函数中超参数的取值，包括以下步骤中的至少一个：

在所述申请空间大于显存碎片的最大值的次数大于第二次数阈值的情况下，增大第N+1轮训练对应的α的取值；

在所述重新计算的时长占总训练时长的百分比相较于上一轮模型训练增加的情况下，减少第N+1轮训练对应的γ的取值。

11.根据权利要求5-10任一项所述的方法，其特征在于，所述方法还包括：

计算第N轮训练的耗时参数＝重新计算耗时/原本计算耗时；

基于所述耗时参数运行模拟退火算法，调整第N+1轮训练对应的显存阈值和/或所述目标估值函数中超参数的取值。

12.根据权利要求4-11任一项所述的方法，其特征在于，所述张量的计算代价的获取步骤包括：

从缓存中读取与所述张量具有相同算子和/或相同输入形状的历史张量的历史计算代价；

将读取到的所述历史计算代价确定为所述张量对应的计算代价。

13.根据权利要求4-12任一项所述的方法，其特征在于，所述张量占用显存的时长的获取步骤包括：

获取当前轮训练中正在执行的算子个数和所述张量进入显存的时间；

根据正在执行的算子个数和所述张量进入显存的时间，确定所述张量占用显存的时长。

14.根据权利要求1-13任一项所述的方法，其特征在于，所述方法还包括：

确定当前轮训练是否包含目标算子执行序列，若是，则将所述目标算子执行序列对应的张量确定为所述目标张量；

其中，所述目标算子执行序列中包括：按照特定顺序执行的多个算子，所述目标算子执行序列对应的张量在历史训练过程中有被释放的记录。

15.根据权利要求1-14任一项所述的方法，其特征在于，在所述释放所述目标张量占用的显存的步骤之后，还包括：

将所述当前轮训练中新产生的张量存储至所释放的显存中。

16.一种模型训练方法，其特征在于，所述方法包括：

获取训练样本集，其中，所述训练样本集中包括用于模型训练的训练数据；

基于所述训练样本集和初始模型进行模型训练，在模型训练过程中，基于权利要求1-15任一项所述的显存管理方法对训练产生的张量进行管理，直至训练得到目标模型。

17.一种显存管理装置，其特征在于，所述装置包括：

第一获取模块，用于获取模型的当前轮训练对应的显存阈值；

第一确定模块，用于在所述电子设备的显存占用值大于所述显存阈值的情况下，确定满足张量选择规则的目标张量；

释放模块，用于释放所述目标张量占用的显存。

18.一种模型训练装置，其特征在于，所述装置包括：

第二获取模块，用于获取训练样本集，其中，所述训练样本集中包括用于模型训练的训练数据；

训练模块，用于基于所述训练样本集和初始模型进行模型训练，在模型训练过程中，基于权利要求17所述的显存管理装置对训练产生的张量进行管理，直至训练得到目标模型。

19.一种电子设备，其特征在于，包括：存储器、处理器及存储在所述存储器上并可在所述处理器上运行的程序，所述程序被所述处理器执行时实现如权利要求1-15任一项所述的显存管理方法的步骤或者实现如权利要求16所述的模型训练方法的步骤。

20.一种计算机可读存储介质，所述计算机可读存储介质上存储有程序，所述程序被处理器执行时实现如权利要求1-15任一项所述的显存管理方法的步骤或者实现如权利要求16所述的模型训练方法的步骤。

Images