基于Tensor访问的深度学习内存管理方法及系统与流程

基于tensor访问的深度学习内存管理方法及系统
技术领域
[0001]
本发明涉及计算机科学人工智能领域，尤其涉及一种针对基于tensor选择的深度学习内存管理方法及系统。


背景技术：

[0002]
深度学习技术的革新大大推动了计算机视觉，自然语言处理，医学等领域的发展。越来越宽的网络结构在训练时具有较高的准确性，因此深度学习社区采用大宽网络已经成为一个趋势。然而以gpu为代表的深度学习加速器的存储容量十分有限，无法容纳深度学习模型在训练过程中的大量数据，严重制约了深度学习技术的发展。越大越深的网络结构需要越多的内存占用，解决深度学习训练的内存不足问题具有重要的价值。现有的解决内存不足主要有数据交换、重新计算、压缩等方法。数据交换是指在gpu训练过程中把暂时不用的数据转出到外部内存中，然后在需要时再转入gpu内存中。重新计算则是将暂时不用的数据释放，然后在需要时重新计算生成。压缩则是将数据进行压缩处理暂存内存，在需要时进行解压缩操作。面对深度学习中产生的众多数据，已有的深度学习内存管理方案往往利用单一手段和局部最优策略降低内存的占有量，没有充分考虑到深度学习的训练性能，限制了深度学习任务的性能提升。


技术实现要素：

[0003]
为了解决深度学习训练过程中内存不足问题的同时最优化深度学习的性能，本发明提出了一种针对于深度学习中数据单元-tensor的选择策略。该选择策略基于整体训练性能的角度设计整数线性规划模型对tensor确定决策。即通过决定深度学习中tensor做某种决策，从而解决内存不足的问题同时获得高效的深度学习训练性能。
[0004]
本发明采用的技术方案是：一种基于tensor访问的深度学习内存管理方法，包含以下几个步骤:步骤1：针对神经网络，收集执行流程中一次迭代过程的tensor访问序列、tensor的大小信息描绘相关决策产生的内存空间开销；同时收集数据交换和重新计算时使用的硬件平台性能信息，包括gpu和cpu之间转移的带宽、gpu上tensor的重新计算时间，并描绘相关决策产生的时间开销；所述相关决策是对tensor访问序列的执行决策，包括：无、转移、释放和重新计算；步骤2：使用步骤1获得的内存空间开销和时间开销信息建立整数线性规划模型，并求解决策序列；步骤3：将步骤2得到的决策序列应用于神经网络之后的每次迭代训练过程中。
[0005]
进一步地，所述步骤1中，描绘相关决策产生的内存空间开销具体为：mem(t) =
ʃ
(s(t,i) = 1) * size(i)其中，mem(t)表示每个阶段t，gpu中tensor的占有量；s(t,i) 表示tensor(i)在t阶段的状态，s的取值为{0,1,2}, 分别表示“不存在”、“在gpu中”，“在cpu中”。i表示tensor编
号，t表示tensor访问序列中的第t个阶段，size(i)为tensor(i)的大小。
[0006]
进一步地，所述步骤1中，描绘相关决策产生的时间开销具体包括：重新计算的时间开销表示为overhead(recomp,t,i)；数据转移引入的时间开销表示为 overhead(swap,t,i) = max{0, cost(i)
ꢀ–ꢀ
time(t-pre)}其中，i表示tensor编号，t表示tensor访问序列中第t个阶段；cost(i)为tensor(i)从gpu或cpu转移的成本：cost(i) = size(i)/bandwidth；size(i)为tensor(i)的大小；time（t-pre）表示pre阶段到t阶段的时间。
[0007]
进一步地，所述步骤2具体为：建立整数线性规划模型：其中，a(i,t)表示tensor在t阶段的决策变量，a的取值为{0,1,2,3,4}, 分别表示“无”、“从gpu转移到cpu”，“从cpu转移到gpu”，“释放”和“重新计算生成”五种决策；利用整数线性规划解决工具来求解上述整数线性规划模型，其中，求解时需要满足的约束条件为：a.在每个训练阶段t，gpu中tensor的占有量小于gpu总内存容量限制；b. 如果在t阶段要访问tensor(i)，那么在t阶段tensor(i)必须在gpu中，即要求s(t,i) = 1；s(t,i) 表示tensor(i)在t阶段的状态，s的取值为{0,1,2}, 分别表示“不存在”、“在gpu中”，“在cpu中”；c. 如果在t阶段tensor(i)需要重新计算生成时，那么它的前序tensor集合j也必须在gpu中，即对于任意 j∈ j，满足s(t,j)=1；进一步地，将决策序列应用于神经网络之后的每次迭代训练过程中，具体为：在前向传播的过程中，在每个阶段t开始时，根据求解的决策方案执行：若tensor(i)需要转移，则执行转出的操作，将tensor(i)的内存地址设为新的地址，同时更新状态；若tensor(i)需要重新计算，则释放tensor(i)的内存，并且记录下生成该tensor(i)的计算操作和输入数据信息，更新状态；在后向传播中，先对所有转出的tensor的再次访问时间进行从短到长排序，当内存足够时开始转入数据；并且在每个后向操作每个阶段开始时执行决策方案：若该阶段中，tensor(i)需要重新计算，则根据生成该tensor(i)的计算操作和输入数据重新执行计算操作。
[0008]
基于上述方法，本发明还提供了一种的内存管理系统，该系统包括：profile模块，用于：收集数据交换和重新计算时使用的硬件平台性能信息，包括gpu和cpu之间转移的带宽、gpu上tensor的重新计算时间等信息，并描绘相关决策产生的时间开销；收集执行流程中一次迭代过程的tensor访问序列、tensor的大小信息描绘相关决策产生的内存空间开销。
[0009]
决策模块，用于将profile模块收集的信息作为整数线性规划模型的输入数据，利
用gurobi ilp求解工具求解出tensor的决策序列；其中，整数线性规划模型具体为： 其中，a(i,t)表示决策变量，a的取值为{0,1,2,3,4}, 分别表示“无”、“从gpu转移到cpu”，“从cpu转移到gpu”，“释放”和“重新计算生成”五种决策；求解时需要满足的约束条件：a. 必须保证在每个训练阶段t，gpu中tensor的占有量小于gpu总内存容量限制；b. 如果在t阶段要访问tensor(i)，那么在t阶段tensor(i)必须在gpu中，即要求s(t,i) = 1；s(t,i) 表示tensor(i)在t阶段的状态，s的取值为{0,1,2}, 分别表示“不存在”、“在gpu中”，“在cpu中”；c. 如果在t阶段tensor(i)需要重新计算生成时，那么它的前序tensor集合j也必须在gpu中，即对于任意 j∈ j，满足s(t,j)=1；overhead(recomp,t,i)表示重新计算的时间开销；overhead(swap,t,i) = max{0, cost(i)
ꢀ–ꢀ
time(t-pre)}表示数据转移引入的时间开销；其中，i表示tensor编号，t表示tensor访问序列中第t个阶段；cost(i)为tensor(i)从gpu和cpu转移的成本：cost(i) = size(i)/bandwidth；size(i)为tensor(i)的大小；time(t-pre)表示pre阶段到t阶段的时间。
[0010]
执行模块，用于将决策模块求解的决策序列应用于神经网络之后的每次迭代训练过程中。
[0011]
进一步地，所述profile模块中，通过在计算操作的前后加时间戳获取某一个计算操作的时间。
[0012]
本发明的有益效果是：本发明方法通过对深度学习中tensor做某种决策，解决神经网络训练时的内存不足问题，实现了对神经网络训练过程中内存的有效管理。相比于现有技术，相同的硬件性能下，本发明可以实现更大的batchsize的神经网络训练。
附图说明
[0013]
图1为本发明方法架构流程图；图2为本发明对tensor序列访问的时间划分为若干个阶段的示意图；图3为具体实施例中在vgg16网络上执行现有capuchin深度学习内存管理方法和本发明方法的对比结果图。
具体实施方式
[0014]
数据交换和数据重新计算是解决内存溢出的常用方法，但现有的单一解决方法不能最优化任务性能，因此我们考虑数据交换和重新计算相结合的方法，而何时对哪些数据执行数据交换、重新计算操作的不同决策对应不同执行流程，并产生不同的内存管理效果。现有的方法使用局部策略执行tensor选择，无法达到好的训练性能。因此，本发明采用整数线性规划确定tensor选择使得内存管理的效果最优，其主要思想是将内存系统中的数据选择决策问题形式化为一个线性规划问题。
[0015]
图1是本发明方法架构流程图，主要为采集硬件相关的成本模型和根据原训练执行流程获得的tensor访问信息、整数线性规划建立和求解、重建带有数据交换和重新计算的训练流程并执行训练。
[0016]
具体来说，包含以下几个步骤:步骤1：针对神经网络，收集执行流程中一次迭代过程的tensor访问序列、tensor的大小信息描绘相关决策产生的内存空间开销；同时收集数据交换和重新计算时使用的硬件平台性能信息，包括gpu和cpu之间转移的带宽、gpu上tensor的重新计算时间等信息，并描绘相关决策产生的时间开销；步骤2：使用步骤1获得的内存空间开销和时间开销信息建立整数线性规划模型，并求解决策序列。
[0017]
步骤3：将步骤2得到的决策序列应用于神经网络之后的每次迭代训练过程中。
[0018]
其中，整数线性规划模型建立和求解是本发明最核心的部分，主要包括决策行为形式化、开销模型建立、约束条件建立、最小化目标函数建立这4个环节。
[0019]
（1）决策时间、决策行为和tensor状态形式化在建立线性规划模型中，为了描述这个开销，本发明对tensor序列访问的时间划分为若干个阶段(stage)。如图2所示，该例子中某次迭代中对tensor的访问顺序依次为tensor 1到tensor 4, 我们就将访问时间划分为4个阶段(stage1~stage4)。
[0020]
当内存不足时，训练执行流程会因为内存溢出而在某个tensor访问时终止。引入决策变量a(i,t)来表示对tensor执行某种决策，其中i表示tensor编号，t表示tensor访问序列中第t个阶段。a的取值为{0,1,2,3,4}, 分别表示“无”、“从gpu转移到cpu”，“从cpu转移到gpu”，“释放”和“重新计算生成”五种决策。
[0021]
与此同时，决策的触发将导致tensor的状态发生变化。引入状态变量s(i,t)来表示tensor(i)在t阶段的状态。s的取值为{0,1,2}, 分别表示“不存在”、“在gpu中”，“在cpu中”。
[0022]
（2）开销模型建立由于数据交换和重新计算的引入，在原来的执行流程上会引入额外的时间开销。从这两种方式来建立开销模型如下：重新计算引入的时间开销重新计算与原计算流程无法并行，因此会在原计算流程中引入额外开销，即计算tensor所需要的时间。在t阶段，当tensor(i)不在内存而需要重新计算tensor(i)时，要求所需前序tensor都在内存中，此时重新计算时间开销表示为overhead(recomp,t,i)。
[0023] 数据转移引入的时间开销数据转移可以和原来的计算流程并行。数据转移带来的额外时间开销由数据转移的固有时间开销和重叠计算时间得到。定义cost(i)为tensor(i)从gpu或cpu转移的成本， 为tensor(i)的大小，即内存占有量，那么：cost(i) = size(i)/bandwidth
ꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀꢀ
（1）假设系统在第pre个阶段开始将数据从gpu转移到cpu中（为t阶段tensor腾出空间），那么该转移带来的额外开销为
overhead(swap,t,i) = max{0, cost(i)
ꢀ–ꢀ
time(t-pre)}
ꢀꢀ
(2)其中cost(t-pre)表示pre阶段到t阶段的时间。当系统在pre阶段开始从cpu朝gpu转移t阶段需要的tensor时，数据转移带来的额外开销也可以类比上述公式进行计算。
[0024]
（3）约束条件建立本发明的目标是在有限内存容量下最小化训练的执行时间，即求解一组a(t,i)使得全局额外时间开销最小。需要满足如下限制条件：a. 必须保证在每个训练阶段t，gpu中tensor的占有量即mem(t)小于gpu总内存容量限制；b. 如果在t时刻要访问tensor，那么在t阶段它必须在gpu中，即要求s(t,i) = 1；c. 如果在t阶段tensor(i)需要重新计算生成时，那么它的前序tensor集合j也必须在gpu中，即对于任意 j∈ j，满足s(t,j)=1。
[0025]
（4）最小化目标函数建立根据上述开销模型和步骤(3)中约束条件，可以得到如下目标函数:（3）最后，利用现成的整数线性规划解决工具来求解获得tensor调度的最优决策序列。
[0026]
基于上述方法，本发明提出了一种内存管理系统，该系统主要包括：profile模块、决策模块和执行模块；该系统可直接添加在深度学习框架（pytorch）上。下面详细阐述这三个模块在系统流程中的实现。
[0027]
1.profile模块profile模块收集中间数据的访问信息和硬件相关的信息。在用户编写完模型代码后，采用profile模块后就可以根据收集的硬件相关的信息获取tensor转移的时间，计算操作执行的时间；具体地，在计算操作的前后加时间戳可以计算某一个计算操作的时间，分别对所有中间数据计算从gpu到cpu转移的时间和从cpu到gpu转移的时间。profile也同时收集神经网络相关信息，包括网络的结构，tensor的数据大小，数据依赖关系。具体地，在程序一个iteration运行的过程中记录下每一个阶段的计算操作、数据大小和计算操作的输入和输出。根据计算操作的输入数据和输出数据，分析数据之间的依赖关系，例如，数据与前序数据的关系。
[0028]
2.决策模块决策模块利用profile模块收集到信息进行决策。首先建立整数线性规划模型，然后将profile模块收集的信息作为模型的输入数据，利用gurobi ilp求解工具求解出tensor的决策序列。最后系统将决策模块中的决策添加到已有的执行流程中。
[0029]
3.执行模块执行模块在原来的pytorch框架上添加了转移和重新计算操作。首先由决策确定了tensor的选择策略。在前向传播的过程中，执行模块在每个操作阶段t开始都会进行判断并根据求解的决策序列执行。如果tensor(i)需要转移，那么执行模块执行转出的操作，将数据的内存地址设为新的地址，同时更新状态；如果tensor(i)需要重新计算，那么执行模块就会释放tensor(i)的内存，并且记录下生成tensor(i)的计算和输入数据信息。在后向传播中，执行模块会对选择的tensor执行转入和重新生成的操作。在后向传播开始时，首先会
对所有转出的tensor的再次访问时间进行排序从短到长，然后当内存足够时就会开始转入数据。另外执行模块在每个后向操作开始时进行判断，如果后向操作计算需要的tensor(i)已经释放了，那么根据生成tensor(i)的操作和输入数据重新执行这个计算操作并更新状态。
[0030]
实施实例在一台配置tesla v100gpu，显存为32gb。cpu型号为intel
®ꢀ
xeon
®ꢀ
gold 6126 cpu@2.60ghz, 操作系统为 ubuntu 18.04.3 lts,cuda toolkit版本为9.0，pytorch的版本为1.5的机器上进行现有的capuchin深度学习内存管理方法和本发明方法的比较。
[0031]
在vgg16网络上执行两种优化方法的训练，结果如图3所示，从训练速度上可以发现，本发明方法在所有的batchsize情况下效果都比capuchin好。在内存占有量的优化上，本发明方法最大支持的batch size为5500，而capuchin的最大的batch size只有4000。
[0032]
显然，上述实施例仅仅是为清楚地说明所作的举例，而并非对实施方式的限定。对于所属领域的普通技术人员来说，在上述说明的基础上还可以做出其他不同形式的变化或变动。这里无需也无法把所有的实施方式予以穷举。而由此所引申出的显而易见的变化或变动仍处于本发明的保护范围。