阿里妹導讀：被譽為國際頂級的機器學習和神經網絡學術會議NIPS 2017于12月4日-9日在美國加州長灘市舉行。在本屆會議上，阿里巴巴除有兩篇論文入選Workshop并進行Oral和Poster形式報告外，三大技術事業(yè)部連續(xù)3天在阿里展區(qū)舉行多場技術研討會，向5000余名參會人員介紹阿里在機器學習、人工智能領域的技術研究、產品與落地應用。

NIPS 2017阿里巴巴展臺-阿里巴巴iDST院長金榕在做演講

這篇介紹深度模型訓練GPU顯存優(yōu)化的論文《TrainingDeeper Models by GPU Memory Optimization on TensorFlow》在NIPS 2017 ML Systems Workshop 中由作者做口頭報告。這篇論文聚焦特征圖，提出兩種方法減少深度神經網絡訓練過程中的顯存消耗，并且把這些方法的實現(xiàn)無縫整合到TensorFlow中，克服了TensorFlow訓練大模型時無法有效優(yōu)化顯存的缺點。

近期深度學習在不同應用中發(fā)揮的作用越來越重要。訓練深度學習模型的必要邏輯包括適合GPU的并行線性代數(shù)計算。但是，由于物理限制，GPU的設備內存（即顯存）通常比主機內存小。最新的高端NVIDIA GPU P100具備12–16 GB的顯存，而一個CPU服務器有128GB的主機內存。然而，深度學習模型的趨勢是「更深更寬」的架構。例如，ResNet 包含多達1001個神經元層，神經網絡機器翻譯（NMT）模型包含8個使用注意力機制的層 ，且NMT模型中的大部分的單個層是按順序水平循環(huán)展開的，難以避免地帶來大量顯存消耗。

簡言之，有限的GPU顯存與不斷增長的模型復雜度之間的差距使顯存優(yōu)化成為必 然。下面將介紹深度學習訓練流程中GPU顯存使用的主要組成。

特征圖（feature map）。對于深度學習模型，特征圖是一個層在前向傳輸中生成的中間輸出結果，且在后向傳輸?shù)奶荻扔嬎阒凶鳛檩斎搿D1是ResNet-50在ImageNet數(shù)據(jù)集上進行一次小批量訓練迭代的GPU顯存占用曲線。隨著特征圖的不斷累積，曲線到達最高點。特征圖的大小通常由批尺寸（batchsize）和模型架構決定（如CNN架構的卷積步幅大小、輸出通道數(shù)量；RNN架構的門數(shù)量、時間步長和隱層大小）。不再需要作為輸入的特征圖占用的顯存將會被釋放，導致圖1中顯存占用曲線的下降。對于復雜的模型訓練，用戶必須通過調整批尺寸，甚至重新設計模型架構來避免「內存不足」的問題。盡管在分布式訓練的情況下 ，訓練任務可以分配到多個設備上來緩解內存不足的問題，但是這也導致了額外的通信開銷。設備的帶寬限制也可能顯著拖慢訓練過程。

圖1：ResNet-50的顯存占用在一個訓練步中的變化曲線。橫軸代表分配／釋放次數(shù)，縱軸代表當前顯存占用的總比特數(shù)。

權重。與特征圖相比，權重占用內存相對較少 。在這篇論文中，權重作為GPU內存中的持久內存，只有整個訓練任務完成后才可以被釋放。

臨時顯存（Temporary memory）。一些算法（如基于Fast-Fourier-Transform（FFT）的卷積算法）需要大量的額外顯存。這些顯存占用是暫時的，在計算結束后立即得到釋放。臨時顯存的大小可以通過在GPU軟件庫（如cuDNN）中列舉每個算法來自動調整，因此可以被忽略。

很明顯，特征圖是GPU顯存使用的主要組成部分。論文作者聚焦特征圖，提出了兩種方法來解決GPU顯存限制問題，即通用的「swap-out/in」方法以及適用于Seq2Seq模型的內存高效注意力層。所有這些優(yōu)化都基于TensorFlow 。TensorFlow具備內置內存分配器，實現(xiàn)了「best-fit with coalescing」的算法。該分配器旨在通過coalescing支持碎片整理（de-fragmentation）。但是，它的內置內存管理策略未考慮大模型訓練時的顯存優(yōu)化。

該論文的貢獻如下。聚焦于特征圖，提出兩種方法減少深度神經網絡訓練過程中的GPU顯存消耗。基于數(shù)據(jù)流圖的「swap-out/in」方法使用主機內存作為更大的內存池，從而放寬GPU顯存上限的限制；而內存高效的注意力層可用來優(yōu)化顯存消耗量大的Seq2Seq模型。這些方法的實現(xiàn)被無縫整合到TensorFlow中，且可透明地應用于所有模型，無需對現(xiàn)有模型架構的描述作任何改變。

作者：孟晨、孫敏敏、楊軍、邱明輝、顧揚

論文地址：
https://github.com/LearningSys/nips17/blob/9ee207c054cf109bc4a068b1064b644d75d0381f/assets/papers/paper_18.pdf

摘要：隨著大數(shù)據(jù)時代的到來、GPGPU的獲取成本降低以及神經網絡建模技術的進步，在GPU上訓練深度學習模型變得越來越流行。然而，由于深度學習模型的內在復雜性和現(xiàn)代GPU的顯存資源限制，訓練深度模型仍然是一個困難的任務，尤其是當模型大小對于單個GPU而言太大的時候。在這篇論文中，我們提出了一種基于通用數(shù)據(jù)流圖的GPU顯存優(yōu)化策略，即「swap-out/in」，將主機內存當做一個更大的內存池來克服GPU的內存限制。同時，為了優(yōu)化內存消耗大的Seq2Seq模型，我們還提出了專用的優(yōu)化策略。我們將這些策略無縫整合到TensorFlow中，且優(yōu)化不會造成準確率的損失。我們在大量的實驗中觀察到了顯著的顯存使用降低。給定一個固定的模型和系統(tǒng)配置，最大訓練批尺寸可以增加2到30倍。

圖2：引用計數(shù)（referencecount）

圖3：swapout/in 優(yōu)化的原子操作（Atomic operation）。刪除從節(jié)點e到節(jié)點b的引用邊，并添加了紅色和藍色的節(jié)點和邊。

圖4：注意力操作（Attentionoperation）優(yōu)化。d指梯度。圖左未經優(yōu)化，圖右經過了顯存優(yōu)化。

表1：對swap out/in 的評估，GPU的顯存上限是12GB。

表2：對顯存高效序列模型的評估。

總結

以上是生活随笔為你收集整理的阿里NIPS 2017论文解读：如何降低TensorFlow训练的显存消耗？的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。