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

权利要求书

1.一种显存管理方法，其特征在于，包括：

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

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

其中，所述第一显存空间与所述第二显存空间在所述显存池中不具有重叠区域；所述第一类信息占用的显存块能够被回收，所述第二类信息占用的显存块不能被回收；每次分配显存块都按照显存块是否可回收，分别在所述第一显存空间和所述第二显存空间中进行。

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

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

3.根据权利要求2所述的方法，其特征在于，所述第一显存空间包括第一队列和第二队列，所述第二显存空间包括第三队列和第四队列；

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

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

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

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

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

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

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

4.根据权利要求3所述的方法，其特征在于，在所述第一类信息的显存分配需求更新的情况下，将分配给第一类信息的显存块回收至所述第一显存空间，包括：

在神经网络模型的反向计算结束时，将分配给第一显存占用信息的显存块回收到第一队列，将分配给第二显存占用信息的显存块回收到第二队列。

5.根据权利要求3所述的方法，其特征在于，所述第一显存占用信息包括神经网络模型反向计算时神经网络层的权重梯度信息；

所述第二显存占用信息包括神经网络模型前向计算时神经网络层的输入信息和输出信息；

所述第三显存占用信息包括神经网络模型的权重信息，和/或，神经网络模型优化器计算时的中间变量；

所述第四显存占用信息包括长短期记忆网络的历史信息，和/或，循环神经网络的历史信息。

6.根据权利要求1所述的方法，其特征在于，包括：

调用显存分配函数，向驱动程序请求获取显存池；

将所述显存池划分为第一显存空间和第二显存空间。

7.根据权利要求1所述的方法，其特征在于，所述第一显存空间的显存块的地址信息是连续的，所述第二显存空间的显存块的地址信息是连续的。

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

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

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

其中，所述第一显存空间与所述第二显存空间在所述显存池中不具有重叠区域；所述第一类信息占用的显存块能够被回收，所述第二类信息占用的显存块不能被回收；每次分配显存块都按照显存块是否可回收，分别在所述第一显存空间和所述第二显存空间中进行。

9.根据权利要求8所述的装置，其特征在于，包括：

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