0
| 本文作者： 奕欣 | 2017-12-18 09:49 |
雷鋒網 AI 科技評論按：本文轉載自“阿里技術”，雷鋒網 AI 科技評論獲授權轉載，并做了不改動原意的編輯。
NIPS 2017 在美國長灘舉辦，場面非常熱烈。在本屆會議上，阿里巴巴除有兩篇論文入選 Workshop 并進行 Oral 和 Poster 形式報告外，三大技術事業部連續 3 天（5 日-7 日）在阿里展區舉行多場技術研討會，向 5000 余名參會人員介紹阿里在機器學習、人工智能領域的技術研究、產品與落地應用。
（NIPS 2017 阿里巴巴展臺-阿里巴巴 iDST 院長金榕進行演講）
這篇介紹深度模型訓練 GPU 顯存優化的論文《Training Deeper Models by GPU Memory Optimization on TensorFlow》將于 8 日在 NIPS 2017 ML Systems Workshop 中由作者做口頭報告。這篇論文聚焦特征圖，提出兩種方法減少深度神經網絡訓練過程中的顯存消耗，并且把這些方法的實現無縫整合到 TensorFlow 中，克服了 TensorFlow 訓練大模型時無法有效優化顯存的缺點。
近期深度學習在不同應用中發揮的作用越來越重要。訓練深度學習模型的必要邏輯包括適合 GPU 的并行線性代數計算。但是，由于物理限制，GPU 的設備內存（即顯存）通常比主機內存小。最新的高端 NVIDIA GPU P100 具備 12–16 GB 的顯存，而一個 CPU 服務器有 128GB 的主機內存。然而，深度學習模型的趨勢是「更深更寬」的架構。例如，ResNet 包含多達 1001 個神經元層，神經網絡機器翻譯（NMT）模型包含 8 個使用注意力機制的層，且 NMT 模型中的大部分的單個層是按順序水平循環展開的，難以避免地帶來大量顯存消耗。
簡言之，有限的 GPU 顯存與不斷增長的模型復雜度之間的差距使顯存優化成為必然。下面將介紹深度學習訓練流程中 GPU 顯存使用的主要組成。
特征圖（feature map）
對于深度學習模型，特征圖是一個層在前向傳輸中生成的中間輸出結果，且在后向傳輸的梯度計算中作為輸入。圖 1 是 ResNet-50 在 ImageNet 數據集上進行一次小批量訓練迭代的 GPU 顯存占用曲線。隨著特征圖的不斷累積，曲線到達最高點。特征圖的大小通常由批尺寸（batch size）和模型架構決定（如 CNN 架構的卷積步幅大小、輸出通道數量；RNN 架構的門數量、時間步長和隱層大小）。不再需要作為輸入的特征圖占用的顯存將會被釋放，導致圖 1 中顯存占用曲線的下降。對于復雜的模型訓練，用戶必須通過調整批尺寸，甚至重新設計模型架構來避免「內存不足」的問題。盡管在分布式訓練的情況下，訓練任務可以分配到多個設備上來緩解內存不足的問題，但是這也導致了額外的通信開銷。設備的帶寬限制也可能顯著拖慢訓練過程。
圖 1：ResNet-50 的顯存占用在一個訓練步中的變化曲線。橫軸代表分配／釋放次數，縱軸代表當前顯存占用的總比特數。
權重
與特征圖相比，權重占用內存相對較少。在這篇論文中，權重作為 GPU 內存中的持久內存，只有整個訓練任務完成后才可以被釋放。
臨時顯存（Temporary memory）
一些算法（如基于 Fast-Fourier-Transform（FFT）的卷積算法）需要大量的額外顯存。這些顯存占用是暫時的，在計算結束后立即得到釋放。臨時顯存的大小可以通過在 GPU 軟件庫（如 cuDNN）中列舉每個算法來自動調整，因此可以被忽略。
很明顯，特征圖是 GPU 顯存使用的主要組成部分。論文作者聚焦特征圖，提出了兩種方法來解決 GPU 顯存限制問題，即通用的「swap-out/in」方法以及適用于 Seq2Seq 模型的內存高效注意力層。所有這些優化都基于 TensorFlow。TensorFlow 具備內置內存分配器，實現了「best-fit with coalescing」的算法。該分配器旨在通過 coalescing 支持碎片整理（de-fragmentation）。但是，它的內置內存管理策略未考慮大模型訓練時的顯存優化。
《Training Deeper Models by GPU Memory Optimization on TensorFlow》的論文貢獻如下：聚焦于特征圖，提出兩種方法減少深度神經網絡訓練過程中的 GPU 顯存消耗。基于數據流圖的「swap-out/in」方法使用主機內存作為更大的內存池，從而放寬 GPU 顯存上限的限制；而內存高效的注意力層可用來優化顯存消耗量大的 Seq2Seq 模型。這些方法的實現被無縫整合到 TensorFlow 中，且可透明地應用于所有模型，無需對現有模型架構的描述作任何改變。
論文： Training Deeper Models by GPU Memory Optimization on TensorFlow
作者：孟晨、孫敏敏、楊軍、邱明輝、顧揚
摘要：隨著大數據時代的到來、GPGPU 的獲取成本降低以及神經網絡建模技術的進步，在 GPU 上訓練深度學習模型變得越來越流行。然而，由于深度學習模型的內在復雜性和現代 GPU 的顯存資源限制，訓練深度模型仍然是一個困難的任務，尤其是當模型大小對于單個 GPU 而言太大的時候。在這篇論文中，我們提出了一種基于通用數據流圖的 GPU 顯存優化策略，即「swap-out/in」，將主機內存當做一個更大的內存池來克服 GPU 的內存限制。同時，為了優化內存消耗大的 Seq2Seq 模型，我們還提出了專用的優化策略。我們將這些策略無縫整合到 TensorFlow 中，且優化不會造成準確率的損失。我們在大量的實驗中觀察到了顯著的顯存使用降低。給定一個固定的模型和系統配置，最大訓練批尺寸可以增加 2 到 30 倍。
圖 2：引用計數（reference count） 
圖 3：swap out/in 優化的原子操作（Atomic operation）
刪除從節點 e 到節點 b 的引用邊，并添加了紅色和藍色的節點和邊。
圖 4：注意力操作（Attention operation）優化。
d 指梯度。圖左未經優化，圖右經過了顯存優化。
表 1：對 swap out/in 的評估。GPU 的顯存上限是 12GB。
表 2：對顯存高效序列模型的評估。
更多資訊，敬請關注雷鋒網 AI 科技評論。
雷峰網版權文章，未經授權禁止轉載。詳情見轉載須知。
