一种基于PyTorch框架的自动重计算方法、装置制造方法及图纸

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

【技术实现步骤摘要】
一种基于PyTorch框架的自动重计算方法、装置
本专利技术涉及网络训练过程重计算
，具体涉及一种基于PyTorch框架的自动重计算方法、装置。
技术介绍
深度学习网络近年来朝着更深、更大的方向发展。在计算机视觉领域，超过1000层的ResNet模型，以及自然语言理解领域的BERT、T5模型，在训练阶段都对GPU的显存具有很高的需求。更深的网络，具有更复杂的函数空间，因此对于复杂的数据场景来说，更深的网络可以更好得拟合数据。而现代GPU支持的显存最大只有32G，由于GPU的价格昂贵，扩展GPU数量对企业成本造成很高的压力。PyTorch中的重计算方式，需要用户手动对网络进行分割，在同一分割内的网络层，其产生的中间值都不需要存储，而只需要存储第一层的输入值。虽然PyTorch的重计算可以有效减小大模型对显存的需求量，但这种手动分割网络的方式对用户理解模型的深度提出了要求。这降低了使用重计算的灵活性，给用户造成了使用门槛。
技术实现思路
针对手动分割网络的方式对用户理解模型的深度提出了要求，降低了使用重计算的灵活性本文档来自技高网...

【技术保护点】
1.一种基于PyTorch框架的自动重计算方法，其特征在于，包括如下步骤：/n在用户定义的深度学习模型中提取计算模块；/n获取计算模块的信息并以获取的信息为结点创建多叉树；/n按照层次遍历的顺序遍历多叉树，逐个选择分割点；将选择的分割点与该层次的其他结点组成分割方案并根据预设的分割点数量选择分割方案；其中，该层次的其他结点为分割点结点所在层次的子结点为非叶结点的结点；/n分割方案选择后，对分割点以及子结点组成的子树进行封装。/n

【技术特征摘要】
1.一种基于PyTorch框架的自动重计算方法，其特征在于，包括如下步骤：
在用户定义的深度学习模型中提取计算模块；
获取计算模块的信息并以获取的信息为结点创建多叉树；
按照层次遍历的顺序遍历多叉树，逐个选择分割点；将选择的分割点与该层次的其他结点组成分割方案并根据预设的分割点数量选择分割方案；其中，该层次的其他结点为分割点结点所在层次的子结点为非叶结点的结点；
分割方案选择后，对分割点以及子结点组成的子树进行封装。


2.根据权利要求1所述的一种基于PyTorch框架的自动重计算方法，其特征在于，获取的计算模块的信息包括模块的名字、从模型索引到该模块的路径、子结点序列、当前结点存储的子模块、输入输出信息。


3.根据权利要求1所述的一种基于PyTorch框架的自动重计算方法，其特征在于，按照层次遍历的顺序遍历多叉树，逐个选择分割点；将选择的分割点与该层次的其他结点组成分割方案并根据预设的分割点数量选择分割方案的步骤包括：
按照层次遍历的顺序遍历多叉树，当遍历到某一层次，将子结点只有叶结点的结点作为分割点并加入必选列表；
将必选列表中的分割点与该层次的其他结点组成分割方案；
根据预设的分割点数量选择分割方案。


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


5.根据权利要求1所述的一种基于PyTorch框架的自动重计算方法，其特征在于，分割方案选择后，对分割点以及子结点组成的子树进行封装的步骤包括：
创建一个继承自PyTorch框架中的模块类；
在创建的模块类内创建一个封装函数，并在创建的封装函数里调用分割点的前向传播函数；
在前向传播函数里，调用PyTorch的检查点函数，并将创建的封装函数传进去完成对分割点的封装。


6.一种基于PyTorch框架的自动重计算装置，其特征在于，...

【专利技术属性】
技术研发人员：王萌，
申请(专利权)人：苏州浪潮智能科技有限公司，
类型：发明
国别省市：江苏;32