[发明专利]一种基于PyTorch框架的自动重计算方法、装置

说明书

本发明提供一种基于PyTorch框架的自动重计算方法、装置，所述方法包括如下步骤：在用户定义的深度学习模型中提取计算模块；获取计算模块的信息并以获取的信息为结点创建多叉树；按照层次遍历的顺序遍历多叉树，逐个选择分割点；将选择的分割点与该层次的其他结点组成分割方案并根据预设的分割点数量选择分割方案；其中，该层次的其他结点为分割点结点所在层次的子结点为非叶结点的结点；分割方案选择后，对分割点以及子结点组成的子树进行封装。采用自动解析模型并分割模型，从而让用户能更有灵活地使用重计算。从而降低了用户的使用门槛。

技术领域

本发明涉及网络训练过程重计算技术领域，具体涉及一种基于PyTorch框架的自动重计算方法、装置。

背景技术

深度学习网络近年来朝着更深、更大的方向发展。在计算机视觉领域，超过1000层的ResNet模型，以及自然语言理解领域的BERT、T5模型，在训练阶段都对GPU的显存具有很高的需求。更深的网络，具有更复杂的函数空间，因此对于复杂的数据场景来说，更深的网络可以更好得拟合数据。而现代GPU支持的显存最大只有32G，由于GPU的价格昂贵，扩展GPU数量对企业成本造成很高的压力。

PyTorch中的重计算方式，需要用户手动对网络进行分割，在同一分割内的网络层，其产生的中间值都不需要存储，而只需要存储第一层的输入值。虽然PyTorch的重计算可以有效减小大模型对显存的需求量，但这种手动分割网络的方式对用户理解模型的深度提出了要求。这降低了使用重计算的灵活性，给用户造成了使用门槛。

发明内容

针对手动分割网络的方式对用户理解模型的深度提出了要求，降低了使用重计算的灵活性，给用户造成了使用门槛的问题，本发明一种基于PyTorch框架的自动重计算方法、装置。

本发明的技术方案是：

一方面，本发明技术方案提供一种基于PyTorch框架的自动重计算方法，包括如下步骤：

在用户定义的深度学习模型中提取计算模块；

获取计算模块的信息并以获取的信息为结点创建多叉树；

按照层次遍历的顺序遍历多叉树，逐个选择分割点；将选择的分割点与该层次的其他结点组成分割方案并根据预设的分割点数量选择分割方案；其中，该层次的其他结点为分割点结点所在层次的子结点为非叶结点的结点；

分割方案选择后，对分割点以及子结点组成的子树进行封装。对PyTorch中的重计算方式进行改进，可自动分割网络，从而降低了用户的使用门槛。

进一步的，获取的计算模块的信息包括模块的名字、从模型索引到该模块的路径、子结点序列、当前结点存储的子模块、输入输出信息。将计算模块结构形成一个多叉树，即，创建的多叉树中的结点包括模块的名字、从模型索引到该模块的路径、子结点序列、当前结点存储的子模块、输入输出信息。

进一步的，按照层次遍历的顺序遍历多叉树，逐个选择分割点；将选择的分割点与该层次的其他结点组成分割方案并根据预设的分割点数量选择分割方案的步骤包括：

按照层次遍历的顺序遍历多叉树，当遍历到某一层次，将子结点只有叶结点的结点作为分割点并加入必选列表；

将必选列表中的分割点与该层次的其他结点组成分割方案；

根据预设的分割点数量选择分割方案。

进一步的，根据预设的分割点数量选择分割方案的步骤包括：

统计分割方案中结点的数量，若分割方案中的结点的数量大于预设的分割点数量，遍历结束，并返回当前的分割方案，否则继续下一层次的遍历。

进一步的，分割方案选择后，对分割点以及子结点组成的子树进行封装的步骤包括：

创建一个继承自PyTorch框架中的模块类；