[发明专利]基于Tensor访问的深度学习内存管理方法及系统

说明书

本发明提供了一种基于Tensor访问的深度学习内存管理方法，该方法通过收集神经网络的执行信息和硬件平台的性能信息获得相关决策下的内存空间开销和时间开销，并建立整数线性规划模型，通过在约束条件下优化求解最优的Tensor调度策略，从而解决内存不足问题的同时获得较高的深度学习训练性能。相比于现有技术，相同的硬件性能下，本发明可以实现更大的batchsize的神经网络训练。本发明同时还提出了一种内存管理系统，包括profile模块、决策模块和执行模块；该系统可直接添加在深度学习框架上，使用方便。

技术领域

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

背景技术

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

发明内容

为了解决深度学习训练过程中内存不足问题的同时最优化深度学习的性能，本发明提出了一种针对于深度学习中数据单元-Tensor的选择策略。该选择策略基于整体训练性能的角度设计整数线性规划模型对tensor确定决策。即通过决定深度学习中Tensor做某种决策，从而解决内存不足的问题同时获得高效的深度学习训练性能。

本发明采用的技术方案是：

一种基于Tensor访问的深度学习内存管理方法，包含以下几个步骤:

步骤1：针对神经网络，收集执行流程中一次迭代过程的Tensor访问序列、Tensor的大小信息描绘相关决策产生的内存空间开销；同时收集数据交换和重新计算时使用的硬件平台性能信息，包括GPU和CPU之间转移的带宽、GPU上Tensor的重新计算时间，并描绘相关决策产生的时间开销；所述相关决策是对Tensor访问序列的执行决策，包括：无、转移、释放和重新计算；

步骤2：使用步骤1获得的内存空间开销和时间开销信息建立整数线性规划模型，并求解决策序列；

步骤3：将步骤2得到的决策序列应用于神经网络之后的每次迭代训练过程中。

进一步地，所述步骤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)的大小。

进一步地，所述步骤1中，描绘相关决策产生的时间开销具体包括：

重新计算的时间开销表示为overhead(recomp,t,i)；

数据转移引入的时间开销表示为 overhead(swap,t,i) = max{0, cost(i) –time(t-pre)}