本文来自AI新媒体量子位(QbitAI)
GPU内存太小可能是神经网络训练过程中最大的拦路虎。
不怕,用这个OpenAI推出的gradient-checkpointing程序包,对于前馈模型来说,仅仅需要增加20%的计算时间,这个程序包,GPU就能适应十倍大的模型。
还有这种操作?
训练神经网络对内存的要求随着网络的深度和batch-size呈线性增长。在内存有限的情况下,如果想训练深层模型,并且增加batch-size,很多研究人员会采用KFAC这样的二阶方法。与小批量的SGD相比,这种方法发需要学习较少的样例。
重点来了。昨天,OpenAI的研究科学家Tim Salimans和前Google Brain工程师的数据科学家Yaroslav Bulatov两人发布了一个python/TensorFlow包,名为gradient-checkpointing。
这个程序包使用了“用亚线性的存储成本训练神经网络”的技术,为简单的前馈网络提供了等价的内存存储,同时能为一般的神经网络节省内存,比如多层架构。
将这个程序包应用到TensorFlow官方CIFAR10 ResNet示例中。在batch size=1280的情况下,将内存和执行时间情况如下图所示。
常规反向传播为线性扩展,但优化后的方法以深度的平方根方式扩展。当我们在更深层次的网络上尝试时,差异就更明显了。
用标准方法,运行这个迭代需要60GB的内存,但新方法只需6GB的RAM。
再来看看计算时间。在实验中,在GTX1080上的运行时间增加了20%,在V100 GPU上时间增加了30%。
如果想了解这个程序包是如何节约内存的,可以移步GitHub一探究竟:
https://github.com/openai/gradient-checkpointing
本文作者:安妮
原文发布时间:2018-01-16