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

权利要求书

1.一种基于PyTorch框架的自动重计算方法，其特征在于，包括如下步骤：

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

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

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

分割方案选择后，对分割点以及子结点组成的子树进行封装；

获取的计算模块的信息包括模块的名字、从模型索引到该模块的路径、子结点序列、当前结点存储的子模块、输入输出信息。

2.根据权利要求1所述的一种基于PyTorch框架的自动重计算方法，其特征在于，按照层次遍历的顺序遍历多叉树，逐个选择分割点；将选择的分割点与该层次的其他结点组成分割方案并根据预设的分割点数量选择分割方案的步骤包括：

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

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

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

3.根据权利要求2所述的一种基于PyTorch框架的自动重计算方法，其特征在于，根据预设的分割点数量选择分割方案的步骤包括：统计分割方案中结点的数量，若分割方案中的结点的数量大于预设的分割点数量，遍历结束，并返回当前的分割方案，否则继续下一层次的遍历。

4.根据权利要求1所述的一种基于PyTorch框架的自动重计算方法，其特征在于，分割方案选择后，对分割点以及子结点组成的子树进行封装的步骤包括：

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

在创建的模块类内创建一个封装函数，并在创建的封装函数里调用分割点的前向传播函数；

在前向传播函数里，调用PyTorch的检查点函数，并将创建的封装函数传进去完成对分割点的封装。

5.一种基于PyTorch框架的自动重计算装置，其特征在于，包括提取模块、创建模块、分割方案生成模块、封装模块；

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

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

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

封装模块，用于分割方案选择后，对分割点以及子结点组成的子树进行封装；

获取的计算模块的信息包括模块的名字、从模型索引到该模块的路径、子结点序列、当前结点存储的子模块、输入输出信息。

6.根据权利要求5所述的一种基于PyTorch框架的自动重计算装置，其特征在于，分割方案生成模块包括遍历单元、分割方案形成单元、确定单元；

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

分割方案形成单元，用于将必选列表中的分割点与该层次的其他结点组成分割方案；

确定单元，用于根据预设的分割点数量选择分割方案。

7.根据权利要求6所述的一种基于PyTorch框架的自动重计算装置，其特征在于，确定单元，具体用于统计分割方案中结点的数量，与预设的分割点数量进行比较；若分割方案中的结点的数量大于预设的分割点数量，输出信息到遍历单元，使遍历结束，并返回当前的分割方案，否则输出信息到遍历单元触发继续下一层次的遍历。

8.根据权利要求7所述的一种基于PyTorch框架的自动重计算装置，其特征在于，封装模块包括类创建单元、函数创建单元、封装单元；类创建单元，用于创建一个继承自PyTorch框架中的模块类；

函数创建单元，用于在创建的模块类内创建一个封装函数，并在创建的封装函数里调用分割点的前向传播函数；

封装单元，用于在前向传播函数里，调用PyTorch的检查点函数，并将创建的封装函数传进去完成对分割点的封装。