[发明专利]显存管理、模型训练方法、装置、电子设备及存储介质

说明书

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

技术领域

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

背景技术

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

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

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

发明内容

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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