OpenAI推新程序包：GPU適應十倍大模型僅需增加20％訓練時間

安妮 編譯自 Medium

量子位 出品 | 公眾號 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

—完—

加入社群

量子位AI社群13群開始招募啦，歡迎對AI感興趣的同學，加小助手微信qbitbot5入群；

此外，量子位專業細分群(自動駕駛、CV、NLP、機器學習等)正在招募，面向正在從事相關領域的工程師及研究人員。

進群請加小助手微信號qbitbot5，並務必備註相應群的關鍵詞~通過審核後我們將邀請進群。（專業群審核較嚴，敬請諒解）

誠摯招聘

量子位正在招募編輯/記者，工作地點在北京中關村。期待有才氣、有熱情的同學加入我們！相關細節，請在量子位公眾號(QbitAI)對話界面，回復「招聘」兩個字。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！