OpenAI开源最新工具包，模型增大10倍只需额外增加20% 计算时间

   <p><a href="/misc/goto?guid=4959012320257333561" title="OpenAI"><img alt="OpenAI开源最新工具包，模型增大10倍只需额外增加20% 计算时间" src="https://simg.open-open.com/show/d1e1227aaffa855a2361c143c00beb99.png" /></a></p>    <p>近日，OpenAI 在 GitHub 上开源最新工具包 gradient-checkpointing，该工具包通过设置梯度检查点（gradient-checkpointing）来节省内存资源。据悉，对于普通的前馈模型，可以在计算时间只增加 20% 的情况下，在 GPU 上训练比之前大十多倍的模型。雷锋网(公众号：雷锋网) AI 研习社将该开源信息编译整理如下：</p>    <p><strong>通过梯度检查点（gradient-checkpointing）来节省内存资源</strong></p>    <p>训练非常深的神经网络需要大量内存，利用 Tim Salimans 和 Yaroslav Bulatov 共同开发的 gradient-checkpointing 包中的工具，可以牺牲计算时间来解决内存过小的问题，让你更好地针对模型进行训练。</p>    <p>对于普通的前馈模型，可以在计算时间只增加 20% 的情况下，在 GPU 上训练比之前大十多倍的模型。</p>    <p>训练深度神经网络时，损失的梯度是在内存密集部分通过反向传播（backpropagation）算法来计算的。在训练模型时定义计算图中的检查点，并在这些检查点之间通过反向传播算法重新计算这些图，可以在降低内存的同时计算梯度值。</p>    <p>当训练一个 n 层的深度前馈神经网络时，可以利用这种方式将内存消耗减少到 O (sqrt (n))，代价是需要执行一个额外的前向传递操作。这个库可以在 Tensorflow 中实现这一功能——使用 Tensorflow graph editor 来自动重写后向传递的计算图。</p>    <p style="text-align:center"><img alt="OpenAI开源最新工具包，模型增大10倍只需额外增加20% 计算时间" src="https://simg.open-open.com/show/d97600a580a3494fb4cfeccae21ec333.jpg" /></p>    <p>图：使用常规的 tf.gradients 函数和使用这种优化内存梯度实现法（memory-optimized gradient implementation）训练不同层数的 ResNet 模型时需要的内存对比</p>    <p><strong>大家现在就可以安装</strong></p>    <blockquote>     <p>pip install tf-nightly-gpu</p>     <p>pip install toposort networkx pytest</p>    </blockquote>    <p>当执行这一程序时，需要保证能找到 CUPTI。</p>    <p>这时可以执行</p>    <blockquote>     <p>export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda/extras/CUPTI/lib64</p>    </blockquote>    <p><strong>使用方法</strong></p>    <p>这个库提供嵌入式功能，能对 tf.gradients 函数进行替换，可以输入如下程序来引入相关函数：</p>    <blockquote>     <p>from memory_saving_gradients import gradients</p>    </blockquote>    <p>大家可以像使用 tf.gradients 函数一样使用 gradients 函数来计算参数损失的梯度。</p>    <p>gradients 函数有一个额外的功能——检查点（checkpoints）。</p>    <p>检查点会对 gradients 函数进行指示——在计算图的前向传播中，图中的哪一部分节点是用户想要检查的点。随后，会在后向传播中重新计算检查点之间的节点。</p>    <p>大家可以为检查点提供一系列张量（gradients (ys,xs,checkpoints=[tensor1,tensor2])），或者可以使用如下几个关键词（'collection'、'memory' 、'speed'）来进行设置。</p>    <p><strong>覆盖 tf.gradients 函数</strong></p>    <p>使用 gradients 函数的另一个方法是直接覆盖 tf.gradients 函数，方法如下：</p>    <blockquote>     <p>import tensorflow as tf</p>     <p>import memory_saving_gradients</p>     <p># monkey patch tf.gradients to point to our custom version, with automatic checkpoint selection</p>     <p>def gradients_memory (ys, xs, grad_ys=None, **kwargs):</p>     <p>return memory_saving_gradients.gradients (ys, xs, grad_ys, checkpoints='memory'， **kwargs)</p>     <p>tf.__dict__["gradients"] = gradients_memory</p>    </blockquote>    <p>这样操作之后，所有调用 tf.gradients 函数的请求都会使用新的节省内存的方法。</p>    <p><strong>测试</strong></p>    <p>在测试文件夹中，有已经写好的用于测试代码准确性和不同模型占用内存的脚本。</p>    <p>大家可以执行 ./run_all_tests.sh 来修改代码，并着手测试。</p>    <p style="text-align:center"><img alt="OpenAI开源最新工具包，模型增大10倍只需额外增加20% 计算时间" src="https://simg.open-open.com/show/fc3c6d962627e181df0dfc02213f6b2d.jpg" /></p>    <p>图：在 CIFAR10 数据集上，使用常规的梯度函数和使用最新的优化内存函数，在不同层数的 ResNet 网络下的内存占用情况和执行时间的对比</p>    <p>来自: <a href="/misc/goto?guid=4959012320375405108" id="link_source2">雷锋网</a></p>