阿里NIPS 2017論文解讀：如何降低TensorFlow訓練的視訊記憶體消耗？
阿里妹導讀：被譽為國際頂級的機器學習和神經網路學術會議NIPS 2017於12月4日-9日在美國加州長灘市舉行。在本屆會議上，阿里巴巴除有兩篇論文入選Workshop並進行Oral和Poster形式報告外，三大技術事業部連續3天在阿里展區舉行多場技術研討會，向5000餘名參會人員介紹阿里在機器學習、人工智慧領域的技術研究、產品與落地應用。
NIPS 2017阿里巴巴展臺-阿里巴巴iDST院長金榕在做演講
這篇介紹深度模型訓練GPU視訊記憶體優化的論文《TrainingDeeper Models by GPU Memory Optimization on TensorFlow》在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視訊記憶體佔用曲線。隨著特徵圖的不斷累積，曲線到達最高點。特徵圖的大小通常由批尺寸（batchsize）和模型架構決定（如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）。但是，它的內建記憶體管理策略未考慮大模型訓練時的視訊記憶體優化。
該論文的貢獻如下。聚焦於特徵圖，提出兩種方法減少深度神經網路訓練過程中的GPU視訊記憶體消耗。基於資料流圖的「swap-out/in」方法使用主機記憶體作為更大的記憶體池，從而放寬GPU視訊記憶體上限的限制；而記憶體高效的注意力層可用來優化視訊記憶體消耗量大的Seq2Seq模型。這些方法的實現被無縫整合到TensorFlow中，且可透明地應用於所有模型，無需對現有模型架構的描述作任何改變。
論文： Training Deeper Models by GPUMemory Optimization on TensorFlow
作者：孟晨、孫敏敏、楊軍、邱明輝、顧揚
論文地址：
https://github.com/LearningSys/nips17/blob/9ee207c054cf109bc4a068b1064b644d75d0381f/assets/papers/paper_18.pdf
摘要：隨著大資料時代的到來、GPGPU的獲取成本降低以及神經網路建模技術的進步，在GPU上訓練深度學習模型變得越來越流行。然而，由於深度學習模型的內在複雜性和現代GPU的視訊記憶體資源限制，訓練深度模型仍然是一個困難的任務，尤其是當模型大小對於單個GPU而言太大的時候。在這篇論文中，我們提出了一種基於通用資料流圖的GPU視訊記憶體優化策略，即「swap-out/in」，將主機記憶體當做一個更大的記憶體池來克服GPU的記憶體限制。同時，為了優化記憶體消耗大的Seq2Seq模型，我們還提出了專用的優化策略。我們將這些策略無縫整合到TensorFlow中，且優化不會造成準確率的損失。我們在大量的實驗中觀察到了顯著的視訊記憶體使用降低。給定一個固定的模型和系統配置，最大訓練批尺寸可以增加2到30倍。
圖2：引用計數（referencecount）
圖3：swapout/in 優化的原子操作（Atomic operation）。刪除從節點e到節點b的引用邊，並添加了紅色和藍色的節點和邊。
圖4：注意力操作（Attentionoperation）優化。d指梯度。圖左未經優化，圖右經過了視訊記憶體優化。
表1：對swap out/in 的評估，GPU的視訊記憶體上限是12GB。
表2：對視訊記憶體高效序列模型的評估。
你可能還喜歡
點選下方圖片即可閱讀