[发明专利]显存管理方法、装置、电子设备和计算机可读存储介质

说明书

本申请公开了装置、电子设备和计算机可读存储介质，涉及模型训练领域。具体实现方案为：在显存池的第一显存空间中，为神经网络模型的第一类信息分配显存块；在显存池的第二显存空间中，为神经网络模型的第二类信息分配显存块；其中，第一显存空间与第二显存空间在显存池中不具有重叠区域；第一类信息占用的显存块能够被回收，第二类信息占用的显存块不能被回收。通过将可回收显存块和不可回收显存块存放在不同的显存空间，可以避免产生分散的显存碎片，进而达到可以充分利用显存资源的效果。

技术领域

本申请涉及一种数据处理领域，尤其涉及一种模型训练领域。

背景技术

随着人工智能技术的不断成熟和推广，人工智能的应用逐渐渗入当今社会的各个领域场景。语音识别领域的语音助手，计算机视觉应用中的人脸识别，自然语言处理的机器翻译等多方面给我们带来便利，这依靠的是深度学习技术。深度学习的过程将大规模的数据在多层神经网络上不断的训练，从多层的训练中不断得到抽象程度更高的数据特征，便于复杂问题的解决。专用图形处理器(Graphics Processing Unit,GPU)是当前深度学习平台主要依托的计算单元。在GPU上训练深度学习模型是深度学习技术的支撑。为了达到更高的模型精度，深度学习的数据规模不断扩大、模型的复杂性和模型深度也在持续增加，但现有GPU显存资源限制成为进一步训练大规模任务的瓶颈。因此，有限的GPU显存空间与不断增长的模型复杂度之间的差距使显存优化成为必然。目前，显存优化的一种做法是在模型训练过程中，回收不需要使用的显存块。显存块的回收会产生显存碎片，导致显存资源得不到充分的利用。

发明内容

本申请实施例提供一种显存管理方法、装置、电子设备和计算机可读存储介质，以解决相关技术存在的问题，技术方案如下：

第一方面，本申请实施例提供了一种显存管理方法，包括：

在显存池的第一显存空间中，为神经网络模型的第一类信息分配显存块；

在显存池的第二显存空间中，为神经网络模型的第二类信息分配显存块；

其中，第一显存空间与第二显存空间在显存池中不具有重叠区域；第一类信息占用的显存块能够被回收，第二类信息占用的显存块不能被回收。

上述技术方案，显存池包括第一显存空间和第二显存空间，第一显存空间和第二显存空间互相独立，在显存池中不具有重叠区域。第一显存空间的显存块分配给神经网络模型中的第一类信息，可以被回收。第二显存空间的显存块分配给神经网络模型中的第二类信息，不会被回收。通过将可回收显存块和不可回收显存块存放在不同的显存空间，可以避免产生分散的显存碎片，进而达到可以充分利用显存资源的效果。

在一种实施方式中，该方法还包括：

在第一类信息的显存分配需求更新的情况下，将分配给第一类信息的显存块回收至第一显存空间。

上述技术方案，可回收显存块的分配和回收发生在独立出来的第一显存空间中，使得对于神经网络模型的多次显存分配需求，显存池都可以提供连续的较大的显存块，避免产生显存碎片，达到充分利用显存资源的效果。

在一种实施方式中，第一显存空间包括第一队列和第二队列，第二显存空间包括第三队列和第四队列；

第一类信息包括在神经网络模型训练前能够计算出显存占用量的第一显存占用信息和在神经网络模型训练前不能计算出显存占用量的第二显存占用信息；第二类信息包括在神经网络模型训练前能够计算出显存占用量的第三显存占用信息和在神经网络模型训练前不能计算出显存占用量的第四显存占用信息；

在显存池的第一显存空间中，为神经网络模型的第一类信息分配显存块，包括：

在神经网络模型训练前，在第一队列中，为第一显存占用信息分配显存块；

在神经网络模型训练时，在第二队列中，为第二显存占用信息分配显存块；