如何设置内存分给显存显存不足

2080Ti其实可以当V100用，功能有点强大。
自深度学习浪潮兴起以来，该模式一直在向更大、更深的方向发展。
2012年，拥有五个卷积层的AlexNet首次在视觉任务上展现出了超强的能力。之后，基础模式开始“深化”:2014年， VGG网达到19层；2015年的ResNet和2017年的DenseNet已经将深度提升到了几百层。
模型大小的增加大大提高了性能。因此， ResNet、DenseNet等。被认为是所有主要视觉任务的基本骨干。但同时，机型的增加也意味着对视频内存的需求变得更高。
为什么GPU内存如此重要？
九年前，辛顿等人率先用两个3GB视频内存的GTX 580 GPU高效训练AlexNet 。此后，对视频内存的需求和模型的大小一直同步增长。想要在游戏中获得好的成绩，想要在实验中超越最先进的效果，想要在工程中契合庞大的商业数据等等。，这一切都离不开记忆的加持。
模型增加一层，视频内存会增加一个点。
在深度学习模型中，占用内存的总是那些特别大的张量，比如各层的权重矩阵、计算出来的张量(激活值)、反向传播所需的张量等。在视觉任务中，中间计算的张量占绝大多数。随着模型越来越深、越来越大，每一层的激活值张量都需要保存在主存中。
以ResNet50为例。在模型的训练中，前向传播中50层的计算结果需要存储在视频存储器中，以便后向传播可以使用这些张量来计算梯度。如果您使用ResNet108 ，您需要的视频内存将是ResNet50的两倍多。视频内存的增加，当然带来了模型效果的提升。另一方面，如果没有足够的视频内存，很多工作将无法完成。
记忆力不够，写论文、玩游戏反复受到约束。
在实验室运行模型和写论文的过程中，内存不足的情况并不少见。一般实验室的显卡都是大家共享的，可能大家手里剩下的不多了。甚至，随着顶级模型变得越来越大，没有人有足够的计算能力和内存来重现最终的实验，更不用说超越它的SOTA结果了。
在这种情况下，学生只有两种选择:向导师申请新的GPU资源，或者缩小模型，做一个Mini实验。前者不一定成功，后者可能有各种不完美。如果能在内存有限的情况下运行顶级模型，做实验和写论文会更容易。
此外，无论是在学校还是在公司，计算能力和内存不足都是很常见的。顶级竞争对手的模式结构可能几乎相同。区别在于谁的模型更大，更能处理复杂的样本。更直观的是，排行榜领跑者的模型可能只差十几层，但正是因为内存限制少于十几层，一些模型才错过了冠军。
内存:算法工程的瓶颈
再举一个常见的例子，企业中的算法工程师有足够的计算能力，内存就没那么重要了。但是，如果只采用并行策略来共享视频内存，可能仍然有足够的视频内存，但是每个GPU的计算负载不足。

文章插图

文章插图
4 V100 ，视频内存满了，但是GPU利用率很低。
即使有V100强大的计算能力，在训练大型机型时也很容易占用16GB的视频内存。但由于批次大小不足，上图中每个V100 GPU的利用率只有20%到30% 。只有不断提高每次迭代的数据吞吐量，才能提高GPU的利用率。
MegEngine:视频内存需要优化。
其实对于深度学习的从业者来说，日常应用中远不止以上三种情况。做深度学习，无论是研究还是工程，都会时不时遇到记忆问题。然而，这个问题的优化非常复杂，需要大量的工程实现来缓解。显然，这样的优化应该由深度学习框架来完成。但是在实际应用中不难发现， TensorFlow和PyTorch似乎并没有提供完美的官方解决方案。
但如果我们看一下新势力，情况可能就不一样了。在最近发布的开源深度学习框架MegEngine 1.4版本中，该框架首次引入了动态图形内存优化技术，大大减少了内存占用问题。
具体来说， MegEngine通过复制和优化ICLR 2021 Spotlight论文《动态张量再物质化》(以下简称DTR)实现了“为更多视频内存而计算” 。有了这项技术的加持，模型的内存占用大大减少，同样的硬件可以训练更大的模型，搭载更大的BatchSize 。这样，学生的小显卡也可以开始训练大型号，工程师的服务器也可以站得更全应用。

文章插图

文章插图
对于原本需要16GB视频内存的机型，优化后使用的视频内存峰值降到了4GB 。
MegEngine是一项内存优化技术，它使1060这样的入门级显卡能够训练只能在2080Ti之前加载的型号。11GB视频内存的2080Ti可以挑战原来只能32GB V100训练的机型。要知道， V100的价格是2080Ti的9倍多。
两行代码，记忆“加倍”
如果需要优化自己的视频内存，可能99%的算法工程师都会放弃。最好的方法是告诉深度学习框架本次训练会分配多少视频内存，剩下的留给框架优化。MegEngine的动态图形内存优化就是基于这个逻辑。
通过两行代码，框架可以完全自动优化视频内存，并在MegEngine中隐藏所有优化逻辑和复杂的工程实现。

文章插图

文章插图
如上图所示，动态计算图中导入了DTR内存优化模块，内存释放阈值配置为5GB 。在训练过程中，由于视频内存已经“翻倍” ，批处理大小可以在翻两番后加载到GPU中。
视频内存扩展带来的收益
在许多情况下，提高视频内存利用率的最显著效果是训练更大的模型。在一定程度上，参数数量越多，效果越好。批量越大，梯度更新方向越准确，模型性能越好。MegEngine开发团队做了很多实验，在提高视频内存利用率的同时，保证了训练的高质量。
最简单的验证方法是不断增加批处理大小，看看显卡能走多远。以下两个表格分别显示了DTR技术在PyTorch和MegEngine上加载或不加载的效果。

文章插图

文章插图
如果不使用动态图形内存优化技术， PyTorch上的模型在一次训练迭代中最多只能处理64个样本， MegEngine最多可以处理100个样本。有了DTR ， PyTorch模型可以在一次迭代中处理140个样本， MegEngine可以尝试处理300个样本。
如果换算成模型大小，再加上MegEngine的动态图形内存优化技术，在相同的GPU和批处理大小下，它可以高效地训练出增加近5倍的模型。
MegEngine动态图形内存优化技术
深度学习模型的内存占用一般分为三个部分:权重矩阵、前向传播中间张量和后向传播梯度矩阵(Adam optimizer) 。
权重矩阵和梯度矩阵占用的内存很难优化，每个模型基本都有固定值。前向传播的中间计算结果不一样:随着Batch Size和模型层数的增加，视频内存必然会增加。如果模型很大，中间的计算结果会占据最重要的内存。

文章插图

文章插图
如上图所示，在正向传播中(第一行从左到右) ，蓝色圆圈表示模型的中间计算结果开始占用视频内存。直到前向传播完成，第一行完全变成蓝圈，之前计算占用的视频内存无法释放。
当反向传播开始时(第二行从右向左) ，通过梯度的计算和应用，保留在主存储器中的正向传播张量可以被释放。
显然，要想减少对视频内存的占用，就必须使用前向传播保存的中间计算结果，这也是MegEngine动态图形内存优化的主要方向。
视频内存的交换计算
动态计算图最直接的方法就是用计算或内存来交换视频内存。因此， MegEngine必须首先决定使用哪种技术。
MegEngine团队通过实验发现，计算时间远小于交换时间。例如，它从视频内存中节省了612.5MB 空，用带宽交换视频内存比用计算交换视频内存慢几十倍到几百倍。

文章插图

文章插图
因此，很明显，动态计算图中也应该使用梯度检查点技术，并且应该用计算来交换视频存储器。
下面是梯度检查点的技术原理。前向传播的第三点是检查点，它将始终存储在视频内存中。第四点可以在计算完成后释放视频内存。如果反向传播中需要第四点的值，可以从第三点重新计算第四点的值。

文章插图

文章插图
虽然一般原理不难理解，但怎么做却相当复杂。MegEngine团队从论文《动态张量再物质化》中学习，对其进行了优化，并在MegEngine中实现。
DTR ，最先进的内存优化技术
DTR是一个完全动态的启发式策略。核心思想是当视频内存超过一定阈值时，动态释放一些合适的张量，直到视频内存降到阈值以下。一般来说，释放张量有三个标准:重新计算张量的代价越小越好；占用的视频内存越大越好；在视频记忆中停留的时间越长越好。

文章插图

文章插图
除了从检查点恢复前向传播结果张量的主要成本外， DTR的额外成本在于找到应该释放的最优张量，即计算上图中张量T的f(t)值。为了减少这部分的计算量， MegEngine还采用了两种运行时优化:
不要考虑小张量，它们不加入候选集。
每次需要释放一个张量时，都会随机采样并遍历几个张量，以节省计算成本。
最难的是项目的实现。
虽然DTR在原则上看起来并不复杂，但真正的问题在于提高可用性，也就是将所有细节隐藏在框架的底层，只为开发人员提供最简单的界面。
在这里，我们使用最简单的计算示例再次遵循框架计算，看看MegEngine如何通过使用动态图的计算历史来恢复和释放张量。

文章插图

文章插图
现在，假设输入中有两个张量A和B ，你想计算a*b和a+b ，但最大内存只能容纳三个张量。当黄框计算c=a+b时，视频内存仍然可以保留张量c ，但在下一步中，当绿框计算d=a*b时，只有c可以被释放， d才能被保存。
不幸的是，在下一步中，灰色框需要得到黄色框的计算结果。不过为了节省视频内存， C已经发布了。因此， MegEngine现在需要做的就是重新运行灰盒的计算图，计算c=a+b ，并加载到视频内存中。显然，这将不可避免地需要释放D的视频内存。
这样，鉴于内存的限制， MegEngine会自动选择合适的张量释放，并在必要时重新计算。如果需要重新计算一个张量的结果，比如上图中的D ，则需要具体的历史计算信息(这里是像a+b这样的计算路径) ，同时需要知道两个输入张量A和b 。
所有这样的历史计算信息都是由MegEngine自动获取保存的， MegEngine的工程师在底层已经用C++进行了处理，用户完全不需要考虑。
struct ComputePath {std::shared_ptr op;SmallVector inputs;SmallVector outputs;double compute_time = 0;} *producer;SmallVector users;size_t ref_cnt = 0;
以上是MegEngine底部用来跟踪计算路径信息的结构。其中op表示生成张量的算子；输入和输出分别表示该算子所需的输入和输出张量；Compute_time表示操作员的实际运行时间。
其实在使用MegEngine的过程中，所有张量都是用Python接口创建的，但是框架会相应地跟踪每个张量的具体信息。每当需要访问张量时，无论张量是否在主存中，没有它都可以立即恢复。所有这些复杂的工程操作和操作逻辑都隐藏在MegEngine C++的底层。

文章插图

文章插图
Python代码将被翻译成C++底层实现， C++代码将通过指针(右边绿色部分)管理显卡内存中的实张量。
幸运的是，这个复杂的操作不需要由算法工程师来完成，而是全部留给MegEngine 。
MegEngine可以做的远不止这些，但大部分都像是动态图形内存优化的技术，帮助用户无形中解决实际问题。2020年3月，开源的MegEngine以肉眼可见的速度快速成长，从静态计算图到动态计算图，再到不断提升的训练能力、移动终端推理性能的优化、动态内存的优化...这或许就是开源的魅力所在。只有不断的优化和创新，才能吸引和满足“挑剔”的开发者。MegEngine的下一个功能会是什么？让我们拭目以待。
【如何设置内存分给显存显存不足】