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