魏曉輝 周博文 李洪亮 徐哲文
摘要: 針對卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練中自動算子選擇算法在較大的顯存壓力下性能下降的問題, 將卸載、 重計算與卷積算子選擇統(tǒng)一建模, 提出一種智能算子選擇算法。該算法權(quán)衡卸載和重計算引入的時間開銷與更快的卷積算子節(jié)省的時間, 尋找卸載、 重計算和卷積算子選擇的調(diào)度, 解決了自動算子選擇算法性能下降的問題. 實(shí)驗(yàn)結(jié)果表明, 該智能算子選擇算法比重計算-自動算子選擇算法縮短了13.53%訓(xùn)練時間, 比已有的卸載/重計算-自動算子選擇算法縮短了4.36%的訓(xùn)練時間.
關(guān)鍵詞: 顯存; 卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練; 卷積算子; 卸載; 重計算
中圖分類號: TP391文獻(xiàn)標(biāo)志碼: A文章編號: 1671-5489(2024)02-0302-09
Memory Optimization Algorithm for Convolutional Neural Networks with Operator Selection
WEI Xiaohui, ZHOU Bowen, LI Hongliang, XU Zhewen
(College of Computer Science and Technology, Jilin University, Changchun 130012, China)
Abstract: Aiming at the problem of? the performance degradation of the automatic operator selection algorithm in convolutional neural network training under high memory pressure, we modelled offloading, recomputing and convolutional operator selecting in a unified manner and proposed an intelligent operator selection algorithm. The algorithm weighed the time overhead introduced by offloading and recomputing against the time saved by faster convolutional operators, found the scheduling of offloading, recomputing and convolutional operator selecting, and solved the performancedegradation problem of the automatic operator selection algorithm. The experimental results? show that the intelligent operator selection algorithm reduces training time by 13.53% over the recomputing-automatic operator selection algorithm and by 4.36% over the existing offloading/recomputing-automatic operator selection algorithm.
Keywords: memory; convolutional neural network training; convolutional operator; offloading; recomputing
由于卷積神經(jīng)網(wǎng)絡(luò)卓越的性能, 越來越多的應(yīng)用使用了卷積神經(jīng)網(wǎng)絡(luò), 如圖像識別、 語義分割等. 為加速卷積計算, 研究人員提出了不同的卷積算子, 它們需要不同的工作空間和卷積耗時. 在訓(xùn)練時, 對于卷積算子的選擇, 主流的訓(xùn)練框架如Tensorflow[1], 會根據(jù)空閑顯存自動選擇合適的卷積算子. 但卷積神經(jīng)網(wǎng)絡(luò)的訓(xùn)練需要保存大量的張量以計算梯度, 訓(xùn)練的顯存壓力隨網(wǎng)絡(luò)規(guī)模的增大而增加. 當(dāng)顯存壓力變大, 空閑顯存匱乏時, 已有的自動算子選擇算法就會選擇工作空間需求少但耗時長的卷積算子, 進(jìn)而延長了訓(xùn)練時間.
目前緩解訓(xùn)練顯存壓力的主流方法是卸載和重計算張量. 卸載方法[2-5]將顯存中暫時不用的張量通過PCIE總線傳輸?shù)饺萘扛蟮腃PU內(nèi)存, 等再次使用時再從CPU的內(nèi)存?zhèn)骰谿PU的顯存, 從而緩解訓(xùn)練的存儲壓力, 運(yùn)行效率受PCIE總線傳輸速度的影響. 重計算方法[6-7]保留部分張量作為檢查點(diǎn), 等使用被釋放的張量時再從檢查點(diǎn)重新計算回來, 運(yùn)行效率受檢查點(diǎn)機(jī)制的影響. 通過卸載和重計算釋放更多的顯存, 本文選擇工作空間大但耗時短的卷積算子, 可有效減少訓(xùn)練時間, 但也會引入更多的卸載和重計算時間開銷. 所以將卸載、 重計算和算子選擇開銷統(tǒng)一建模, 并基于開銷模型進(jìn)行綜合考慮設(shè)計算法已經(jīng)成為目前該領(lǐng)域重要的研究內(nèi)容.
本文針對卷積算子選擇的問題, 進(jìn)行整數(shù)規(guī)劃問題的模型設(shè)計. 在此基礎(chǔ)上, 本文提出智能算子選擇算法, 綜合考慮卸載、 重計算和卷積算子選擇, 使卸載和重計算為更快的卷積算子的工作空間騰出顯存, 進(jìn)而縮短訓(xùn)練時間. 本文使用大規(guī)模求解器Gurobi搜索整數(shù)規(guī)劃模型的結(jié)果作為建模整數(shù)規(guī)劃問題的可行解, 并將求解器結(jié)果轉(zhuǎn)化為顯存管理調(diào)度, 在訓(xùn)練時通過調(diào)度管理顯存. 本文對實(shí)際AI生產(chǎn)環(huán)境下的多種網(wǎng)絡(luò)進(jìn)行測試, 并將智能算子選擇與自動算子選擇進(jìn)行對比. 實(shí)驗(yàn)結(jié)果表明, 本文的MO-SOS(memory optimization with smart operator selection)算法可以顯著提升GPU運(yùn)算性能.
1 研究背景
1.1 自動算子選擇與卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練顯存壓力
由于計算量大, 卷積計算是卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練中最耗時的計算. 為加速卷積計算, 研究人員提出了不同的卷積算子實(shí)現(xiàn)卷積計算, 如快速Fourier變換卷積[8]和Winograd變換卷積[9]等. 它們在加速卷積計算時需額外的空間保存計算過程中的臨時變量, 即神經(jīng)網(wǎng)絡(luò)算子的工作空間. 在訓(xùn)練卷積神經(jīng)網(wǎng)絡(luò)時, 現(xiàn)有的訓(xùn)練框架如Tensorflow, 會根據(jù)空閑顯存自動選擇合適的卷積算子.
近年來, 卷積神經(jīng)網(wǎng)絡(luò)的網(wǎng)絡(luò)規(guī)模越來越大, 使訓(xùn)練的顯存壓力變得越來越嚴(yán)峻. 訓(xùn)練卷積神經(jīng)網(wǎng)絡(luò)是尋找一組網(wǎng)絡(luò)參數(shù)使神經(jīng)網(wǎng)絡(luò)的輸出盡可能接近期望值. 圖1為一個多層卷積神經(jīng)網(wǎng)絡(luò)的訓(xùn)練過程, 包括卷積層CONV、 批標(biāo)準(zhǔn)化層BN、 激活層ReLU和全連接層FC等. 目前尋找網(wǎng)絡(luò)參數(shù)常用的方法是梯度下降法, 需要計算神經(jīng)網(wǎng)絡(luò)的損失(Loss)和參數(shù)的梯度.
Loss是神經(jīng)網(wǎng)絡(luò)輸出與期望值之間的差距, 這種差距可通過計算損失函數(shù)得到, 即圖1中的Loss層. Loss是通過前向傳播計算的, 前向傳播是圖1中實(shí)線表示的方向. 神經(jīng)網(wǎng)絡(luò)訓(xùn)練可通過Loss計算出模型的參數(shù)梯度, 從而使用參數(shù)梯度更新模型參數(shù).
計算梯度是通過后向傳播完成的, 即圖1中虛線表示的方向. 計算梯度使用鏈?zhǔn)椒▌t依次計算出每層的參數(shù)梯度. 計算參數(shù)梯度依賴張量[WTHX]a[WTBX]i, 所以卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練會保存前向傳播的輸出張量. 目前, 卷積神經(jīng)網(wǎng)絡(luò)通過增加層數(shù)以獲得更高準(zhǔn)確率, 所以它們的訓(xùn)練顯存需求也會越來越大, 如Vgg19[10]和Resnet1001等. 當(dāng)顯存壓力較大, 空閑顯存匱乏時, 自動算子選擇算法會選擇工作空間少但耗時較長的卷積算子, 增加訓(xùn)練時間. 但卷積神經(jīng)網(wǎng)絡(luò)的張量在訓(xùn)練中大部分時間內(nèi)都是空閑的, 未被訪問[2], 所以可將張量卸載或重計算減輕訓(xùn)練顯存壓力.
1.2 緩解顯存壓力方案與智能卷積算子選擇動機(jī)
卸載將CPU的內(nèi)存作為GPU的緩存, 暫存訓(xùn)練中的張量. 目前的GPU通常作為CPU的外設(shè)設(shè)備, 它是通過PCIE總線與CPU相連, 它們通過PCIE總線相互傳輸數(shù)據(jù). 由于CPU的內(nèi)存比GPU的顯存容量大, 所以可以把張量卸載到CPU主存, 再次需要張量時預(yù)取回GPU顯存. 目前的主流顯卡如3090,P40等, 都支持計算、 卸載和預(yù)取張量并行. 所以為不阻塞訓(xùn)練中的計算, 需要提前卸載或預(yù)取張量. VDNN(virtualized deep neural networks)[2]在前向傳播中的卷積計算時卸載它的輸入張量, 當(dāng)后向傳播需要已卸載的輸入張量時再提前預(yù)取回. SwapAdvisor[4]通過進(jìn)化計算選擇卸載的張量、 卸載開始時間和預(yù)取開始時間, 進(jìn)一步減少阻塞時間. STR(swap dominated tensor re-generation)[11]發(fā)現(xiàn)以前的工作未考慮到多次預(yù)取張量, 因?yàn)榘岩粋€張量卸載到CPU的內(nèi)存后, CPU的內(nèi)存中便有了該張量的備份, 當(dāng)在訓(xùn)練中多次訪問這個張量時, STR認(rèn)為可以把張量多次預(yù)取回GPU顯存.
隨著GPU計算能力越來越強(qiáng), 張量的卸載和預(yù)取的阻塞可能比計算更耗時, 因此研究人員又提出了另一種緩解顯存壓力的方法——重計算. 重計算保留一部分張量作為檢查點(diǎn), 釋放其他的張量騰出顯存空間, 當(dāng)被釋放的張量再次需要訪問時, 從現(xiàn)有的檢查點(diǎn)開始重新計算. 由于引入了重計算, 訓(xùn)練會引入額外的計算時間, 該方法是一種用時間換取空間的方法, 所以保留不同的檢查點(diǎn)會引入不同的重計算時間. Chen等[6]把網(wǎng)絡(luò)分成多個部分, 每部分只保留其輸入, 對于每部分的其他張量都重新計算. Checkmate[7]不同于文獻(xiàn)[6]的工作, 其探究了在顯存容量的限制下, 哪些張量作為檢查點(diǎn)能引入最小的額外計算時間, 將重計算問題抽象為線性規(guī)劃問題, 用現(xiàn)有的求解器解決重計算問題.
卸載和預(yù)取會導(dǎo)致額外的阻塞時間, 而重計算會引入額外的計算時間. 文獻(xiàn)[12]認(rèn)為卷積計算時間較長, 不適合用于重計算, 所以卸載和預(yù)取卷積計算的輸出張量, 而對于激活、 批標(biāo)準(zhǔn)化的輸出張量, 它們的計算時間相對較少, 適合用于重計算. 文獻(xiàn)[13]研究表明, 以往的卸載和重計算的結(jié)合都是以神經(jīng)網(wǎng)絡(luò)層的粒度管理顯存, 然而一個網(wǎng)絡(luò)層中也會存在多個計算, 所以提出了張量粒度的顯存管理方案. STR認(rèn)為在非線性卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練時, 一個張量會被多次訪問, 所以STR通過多次預(yù)取已卸載的張量減少重計算的時間.
當(dāng)顯存壓力較大, 空閑顯存匱乏時, 可通過卸載和重計算釋放更多的顯存, 為卷積算子增加空閑顯存. 如圖2所示, 它們是Vgg19在Tesla P40平臺上訓(xùn)練時的6個前向傳播的卷積計算, 使用了主流的CUDNN作為卷積算子庫. 不可釋放顯存是當(dāng)前卷積計算依賴的輸入和輸出張量占用的顯存, 而可釋放顯存是暫時不用的張量顯存. 與現(xiàn)有的工作相同, 本文主要對前向傳播的輸出張量進(jìn)行顯存優(yōu)化, 所以圖2只統(tǒng)計了前向傳播的輸出張量. 在訓(xùn)練時, 卸載和重計算可釋放的顯存, 為卷積計算騰出空閑顯存. 但卸載和重計算會引入更多的額外時間開銷, 所以將卸載和重計算引入的時間開銷與更快卷積算子節(jié)省的時間權(quán)衡. 本文拓展了STR的模型, 對卸載、 重計算和卷積算子選擇進(jìn)行統(tǒng)一建模, 尋找卸載、 重計算和卷積算子的調(diào)度, 以減少訓(xùn)練時間.
2 問題模型與算法設(shè)計
2.1 考慮卷積算子選擇的顯存優(yōu)化問題描述
由于深度神經(jīng)網(wǎng)絡(luò)框架如Tensorflow和Pytorch等的使用方便與性能強(qiáng)大, 神經(jīng)網(wǎng)絡(luò)設(shè)計者常使用它們設(shè)計和訓(xùn)練網(wǎng)絡(luò). 這些主流的框架將深度神經(jīng)網(wǎng)絡(luò)的訓(xùn)練過程表示成計算圖, 如有向無環(huán)圖(DAG). 計算圖由節(jié)點(diǎn)和邊組成, 節(jié)點(diǎn)表示計算, 如卷積、 矩陣相乘等, 邊表示數(shù)據(jù)依賴, 如圖3中的DAG計算圖. 在訓(xùn)練時, 訓(xùn)練框架會按照計算圖的拓?fù)渑判颍?依次執(zhí)行節(jié)點(diǎn)的計算[4,7].
圖3中v18節(jié)點(diǎn)是Vgg19模型的一個卷積計算, 運(yùn)行在Tesla P40平臺上, 算子庫使用了8.0.5版本的CUDNN. 當(dāng)顯存壓力變大, 空閑顯存匱乏時, 自動算子選擇只能選擇不需要工作空間但耗時最長的算子Implicit GEMM, 增加訓(xùn)練耗費(fèi)的時間. 所以MO-SOS算法要通過卸載和重計算為工作空間釋放顯存, 選擇更快的卷積算子. 選擇最快的算子Winograd Nonfused, 則需要通過卸載和重計算為工作空間釋放約2 GB的顯存. 相應(yīng)地, 卸載和重計算也會引入更多的額外時間開銷. 若引入的額外時間開銷大于最快卷積算子節(jié)省的訓(xùn)練時間, 反而會增加訓(xùn)練時間. 另一種算子的選擇是工作空間少但耗時較長的卷積算子Winograd. Winograd算子只需要通過卸載和重計算為工作空間釋放約6 MB的顯存.
選擇最快或最慢的卷積算子都有可能增加訓(xùn)練時間, 所以需要對卷積算子、 卸載和重計算進(jìn)行統(tǒng)一建模, 尋找一個卷積算子和卸載/重計算的調(diào)度, 使得從v1計算節(jié)點(diǎn)到vn計算節(jié)點(diǎn)執(zhí)行時間最短. 針對卷積算子的選擇問題, 本文從考慮卷積算子選擇問題的定義、 卸載和預(yù)取建模及工作空間的顯存限制三方面進(jìn)行問題模型的設(shè)計.
2.2 考慮卷積算子選擇問題的定義
設(shè)一個訓(xùn)練計算圖為G=(V,E), 其中V={v1,v2,…,vn}表示該計算圖有n個計算節(jié)點(diǎn), E表示計算節(jié)點(diǎn)之間的依賴關(guān)系. 每個計算節(jié)點(diǎn)vi的輸出需要Mi顯存存儲輸出張量. 每個計算節(jié)點(diǎn)vi依賴于其他計算節(jié)點(diǎn)的輸出vj /St+1,j≤∑aj/v=1Rt,j,v+Pss(t,j),j, t,j,(5)其中: ?ak表示k計算節(jié)點(diǎn)有多少種算子實(shí)現(xiàn); O和P變量表示卸載和預(yù)取的決策; 式(2)表示v1到vn計算節(jié)點(diǎn)都要計算一次; 式(3)約束了對于不同的計算節(jié)點(diǎn), 算子的選擇最多只能有一種. 如上所述, 每個計算節(jié)點(diǎn)都依賴于其他計算節(jié)點(diǎn)的輸出張量. 本文使用卸載和重計算釋放顯存空間, 在某個節(jié)點(diǎn)計算時, 它依賴的張量可能并不存在于顯存中, 它可能未重計算或者預(yù)取回顯存. 所以MO-SOS算法使用St,j變量表示在t階段時, j張量是否可以訪問. 式(4)表示在計算每個節(jié)點(diǎn)時它依賴的張量必須可以訪問, DEPS[k]表示k計算節(jié)點(diǎn)的依賴張量. 而式(5)表示如果在t階段重計算或者預(yù)取回了j張量, 則j張量在之后的計算中能被其他計算節(jié)點(diǎn)訪問. 2.3 卸載和預(yù)取建模 由于卸載、 預(yù)取與重計算可以并行, 因此MO-SOS算法可通過卸載和預(yù)取減少需要重計算的張量. 但為保證卸載與預(yù)取的正確性, MO-SOS算法需要對卸載和預(yù)取進(jìn)行約束. 本文規(guī)定在每個節(jié)點(diǎn)計算的開端, 開始卸載或預(yù)取張量. Ot,j表示在t階段是否開始卸載j張量, Pt,j表示在t階段是否開始預(yù)取j張量, 則∑n/t=1Ot,j≤1, j,(6) 式(6)表示每個張量只能卸載一次, 因?yàn)閺埩勘恍遁d到CPU內(nèi)存后, 內(nèi)存中的張量可作為備份多次預(yù)取. 式(7)表示在卸載時張量的數(shù)據(jù)必須是可以訪問的, 否則會傳輸臟數(shù)據(jù). ss(t,i)表示在t階段完整地卸載或預(yù)取i張量, 需要在哪個階段開始卸載或預(yù)??; sf(t,i)表示在t階段開始卸載或預(yù)取i張量, 會在哪個階段完成. 這兩個函數(shù)可以根據(jù)PCIE總線帶寬、 張量大小和每個節(jié)點(diǎn)的計算時間計算得出. 但不同的卷積算子耗費(fèi)不同的卷積時間. 為解決上述問題, 在計算ss和sf函數(shù)時, 可根據(jù)最快的卷積算子計算, 保證卸載和預(yù)取不會阻塞計算. 式(8)表示如果沒有預(yù)取就不會卸載張量, 且m是可設(shè)置的參數(shù). 式(9)表示預(yù)取必須在張量卸載完成后開始. 此外, 本文限制了每個階段只能卸載或預(yù)取一個張量, 即Ot,j≤∑ 2.4 考慮工作空間的顯存限制 為約束每個節(jié)點(diǎn)運(yùn)行時顯存不會溢出, 需要定義一個變量表示每個節(jié)點(diǎn)計算時需要的顯存. 所以用Ut,k表示在t階段計算k節(jié)點(diǎn)時需要的顯存, Ut,0表示在t階段開始時需要的顯存, 它包括模型的輸入、 參數(shù)、 梯度以及作為檢查點(diǎn)的張量Ut,0=Minput+Mparam+Mgrad+∑n/j=1St,jMj+∑(x,y)∈OCC(t)MyPx,y.(14)此外, 預(yù)取也會占用顯存, 所以O(shè)CC(t)表示了哪些張量的預(yù)取在t階段占用顯存. 然后遞歸地計算t階段每個節(jié)點(diǎn)計算時需要的顯存Ut,k:Ut,k+1=Ut,k-mem_free(k)+∑ak+1/v=1Rt,k+1,vMk+1+∑ MO-SOS算法希望通過卸載和重計算騰出更多的顯存空間, 以選擇更快的卷積算子, 進(jìn)而減少訓(xùn)練時間. 式(2)~(5)保證了重計算的正確性, 式(6)~(13)保證了卸載和預(yù)取的正確性, 且式(14)~(18)約束了顯存占用. 該問題是難以解決的NP-難問題, 求解空間巨大, 無法通過高復(fù)雜度的遍歷算法求解, 因此需要有針對性地進(jìn)行模型設(shè)計. 2.5 算法設(shè)計 本文問題可以使用求解器根據(jù)輸入的計算圖和GPU的顯存限制求解一個可行的方案. 目標(biāo)公式(1)中的R,S,O,P, Free都是0~1變量, 所以是一個整數(shù)規(guī)劃問題. MO-SOS算法使用現(xiàn)有的求解器去搜索一個可行的結(jié)果, 例如Gurobi,Cplex,Scip等. Gurobi是大規(guī)模數(shù)學(xué)規(guī)劃優(yōu)化器, 在多個整數(shù)規(guī)劃問題中有更快的優(yōu)化速度, 適合于本文的最小化問題. 根據(jù)Gurobi求解器給出的結(jié)果, MO-SOS算法會產(chǎn)生一個顯存管理調(diào)度L={l1,l2,…,lk}, 其中l(wèi)i是顯存管理的操作, 如計算張量、 釋放張量、 卸載張量和預(yù)取張量. MO-SOS算法如下. MO-SOS算法遍歷所有的R,O,P, Free變量去制定顯存管理調(diào)度. 如果Ri,j,v為1, 則在i階段, MO-SOS算法要計算j張量, 添加顯存管理操作計算張量j到L中, 使用v算子, 即代碼10)~12), Algoj表示j節(jié)點(diǎn)的算子數(shù)量. 如果Oi,k或Pi,k為1, 則在i階段, MO-SOS算法要卸載或預(yù)取k張量, 添加顯存管理操作卸載或預(yù)取k張量到L中, 即代碼4)~6). 如果Freei,j,t為1, 則在i階段, 計算j張量后, MO-SOS算法釋放t張量, 添加顯存管理操作釋放t張量到L中, 即代碼15)~17). 3 實(shí) 驗(yàn) 3.1 實(shí)驗(yàn)環(huán)境 本文選取兩種卷積神經(jīng)網(wǎng)絡(luò)Vgg19和Unet, 使用C++,CUDA和CUDNN實(shí)現(xiàn)Vgg19和Unet的訓(xùn)練, CUDA的版本是11.1, CUDNN的版本是8.0.5. Vgg19在訓(xùn)練時使用數(shù)據(jù)集Imagenet, 圖片大小為227×227, Unet在訓(xùn)練時使用數(shù)據(jù)集Data Science Bowl, 圖片大小為256×256. 它們訓(xùn)練的批次大小分別為144和120. 實(shí)驗(yàn)環(huán)境采用Ubuntu18.04操作系統(tǒng), CPU為Intel(R) Xeon(R) CPU E5-2680 v4, 內(nèi)存為125 GB. 實(shí)驗(yàn)使用的GPU為英偉達(dá)的Tesla P40, 顯存為24 GB, 而PCIE總線的卸載和預(yù)取速度分別至少達(dá)12 GB/s和11 GB/s. 求解整數(shù)規(guī)劃問題的求解器使用Gurobi, 版本為9.5.2. 本文將智能算子選擇算法MO-SOS和自動算子選擇算法Auto進(jìn)行比較. 為能在顯存壓力較大的情況下進(jìn)行實(shí)驗(yàn), 給自動算子選擇算法配置兩種最新的顯存優(yōu)化方法: Checkmate(只有重計算)和STR(卸載和重計算相結(jié)合). 這兩種顯存優(yōu)化算法與自動算子選擇結(jié)合后分別為Checkmate-Auto和STR-Auto算法. 本文設(shè)置了不同的顯存比例, 逐漸減少能使用的顯存, 增加訓(xùn)練的顯存壓力, 測試了Checkmate-Auto,STR-Auto和MO-SOS算法的訓(xùn)練時間. 由于訓(xùn)練是多個訓(xùn)練步組成的, 每個訓(xùn)練步執(zhí)行時間浮動較小, 所以本文對每個神經(jīng)網(wǎng)絡(luò)測試了20個訓(xùn)練步, 取平均值作為每個神經(jīng)網(wǎng)絡(luò)一個訓(xùn)練步的執(zhí)行時間. 本文測試了在不同的顯存壓力下, 智能算子選擇算法和自動算子選擇算法訓(xùn)練步的執(zhí)行時間, 以考察智能算子選擇算法能否縮短訓(xùn)練時間. [HT5H〗3.2 實(shí)驗(yàn)結(jié)果分析 圖4為不同算法對Vgg19和Unet一個訓(xùn)練步的執(zhí)行時間. 由圖4可見, 智能算子選擇算法在多種顯存壓力下都優(yōu)于自動算子選擇算法. 本文的智能算子選擇算法MO-SOS平均比只有重計算的Checkmate-Auto算法減少了13.53%的訓(xùn)練時間, 比卸載和重計算結(jié)合的STR-Auto算法減少了4.36%的訓(xùn)練時間. 在顯存比例為0.8和0.7, 顯存壓力較小時, MO-SOS和STR-Auto算法的訓(xùn)練時間基本相當(dāng), 而結(jié)合了卸載的MO-SOS和STR-Auto算法仍比Checkmate-Auto算法減少約11.2%的訓(xùn)練時間, 這主要是前者使用卸載減少了需要重計算的張量. 隨著顯存壓力變大, MO-SOS算法的訓(xùn)練時間逐漸比STR-Auto算法減少. 這是因?yàn)轱@存壓力變大, STR-Auto算法沒有為更快的卷積算子釋放更多的顯存, 選擇了更慢的卷積算子, 增加了訓(xùn)練時間. 但對于Unet網(wǎng)絡(luò), 在顯存比例為0.4時, MO-SOS與STR-Auto算法的訓(xùn)練時間差距縮小. 不同于Vgg19, Unet是非線性網(wǎng)絡(luò), 前向傳播的輸出張量不僅在后向傳播計算梯度時多次使用, 并且在前向傳播中也會多次使用. 所以Unet在訓(xùn)練時需要長時間保存多個張量, 導(dǎo)致更大的顯存壓力, 難以選擇更快的卷積算子. MO-SOS和STR-Auto算法訓(xùn)練時間比Checkmate-Auto算法少, 這是因?yàn)榍罢咄ㄟ^卸載減少了重計算的張量. 如圖5所示, 在Vgg19訓(xùn)練時, MO-SOS和STR-Auto算法在顯存比例為0.7時只有卸載, 而Checkmate-Auto算法只有重計算. MO-SOS和STR-Auto算法通過卸載騰出了足夠的顯存, 所以它們的卷積計算和其他計算的耗時都明顯少于Checkmate-Auto算法. MO-SOS算法訓(xùn)練時間比STR-Auto算法少, 是因?yàn)镸O-SOS算法卸載和重計算更多的張量以增加空閑顯存, 選擇更快的卷積算子減少訓(xùn)練時間. 由圖5可見, 當(dāng)顯存比例為0.7時, MO-SOS算法比STR-Auto算法卸載的張量更多. 所以MO-SOS算法的卷積計算耗時少于STR-Auto算法. 隨著顯存壓力增大, MO-SOS和STR-Auto算法逐漸增加了重計算的張量以釋放顯存. 當(dāng)顯存比例為0.5時, 為給卷積計算的工作空間騰出更多的顯存, MO-SOS算法卸載和重計算的張量多于STR-Auto算法. 所以MO-SOS算法的其他計算耗時多于STR-Auto算法, 但MO-SOS算法的卷積計算耗時明顯少于STR-Auto算法, 總體訓(xùn)練時間也少于STR-Auto算法. 實(shí)驗(yàn)結(jié)果表明, 可以通過卸載和重計算為卷積計算增加空閑顯存, 選擇更快的卷積算子減少訓(xùn)練時間. 在顯存壓力較小時, 減少的訓(xùn)練時間并不明顯. 但隨著顯存壓力增大, 自動算子選擇算法的空閑顯存逐漸減少, 選擇了更慢的卷積算子. 而智能算子選擇算法通過卸載和重計算釋放了更多的顯存, 選擇了更快的卷積算子, 進(jìn)而縮短了訓(xùn)練時間. 綜上所述, 針對自動算子選擇算法在空閑顯存匱乏時會選擇工作空間需求少但耗時長的卷積算子, 降低了訓(xùn)練的性能、 增加了訓(xùn)練時間的問題, 本文提出了一種智能算子選擇算法, 通過卸載和重計算增加空閑顯存. 本文對卸載、 重計算和卷積算子選擇進(jìn)行了統(tǒng)一的整數(shù)規(guī)劃建模, 使用現(xiàn)有的求解器進(jìn)行求解, 并對多種網(wǎng)絡(luò)在各種顯存壓力下進(jìn)行實(shí)驗(yàn). 實(shí)驗(yàn)結(jié)果表明, 智能算子選擇算法在多種顯存壓力下優(yōu)于自動算子選擇算法. 參考文獻(xiàn) [1]ABADI M, BARHAM P, CHEN J, et al. Tensorflow: A System for Large-Scale Machine Learning [C]//Proceedings of the 12th USENIX Conference on Operating Systems Design and Implementation. New York: ACM, 2016: 265-283. [2]RHU M, GIMELSHEIN N, CLEMONS J, et al. vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design [C]//2016 49th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO). Piscataway, NJ: IEEE, 2016: 1-13. [3]SHRIRAM S B, GARG A, KULKARNI P. Dynamic Memory Management for GPU-Based Training of Deep Neural Networks [C]//2019 IEEE International Parallel and Distributed Processing Symposium (IPDPS). Piscataway, NJ: IEEE, 2019: 200-209. [4]HUANG C C, JIN G, LI J. Swapadvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping [C]//Proceedings of the Twenty-Fifth International Conference on Architectural Support for ProgrammingLanguages and Operating Systems. New York: ACM, 2020: 1341-1355. [5]BAE J, LEE J, JIN Y, et al. FlashNeuron: SSD-Enabled Large-Batch Training of Very Deep Neural Networks [C]//Proceedings of the 12th USENIX Conference on Operating Systems Design and Implementation. New York: ACM, 2021: 387-401. [6]CHEN T Q, XU B, ZHANG C Y, et al. Training Deep Nets with Sublinear Memory Cost [EB/OL]. (2016-04-21)[2023-02-01]. https://arxiv.org/abs/1604.06174. [7]JAIN P, JAIN A, NRUSIMHA A, et al. Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization [EB/OL]. (2019-10-07)[2023-02-01]. https://arxiv.org/abs/1910.02653. [8]MATHIEU M, HENAFF M, LeCUN Y. Fast Training of Convolutional Networks through Ffts [EB/OL]. (2013-12-20)[2023-02-15]. https://arxiv.org/abs/1312.5851. [9]LAVIN A, GRAY S. Fast Algorithms for Convolutional Neural Networks [C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. Piscataway, NJ: IEEE, 2016: 4013-4021. [10]SIMONYAN K, ZISSERMAN A. Very Deep Convolutional Networks for Large-Scale Image Recognition [EB/OL]. (2014-09-04)[2023-01-20]. https://arxiv.org/abs/1409.1556. [11]WEN L J, ZONG Z, LIN L, et al. A Swap Dominated Tensor Re-generation Strategy for Training Deep Learning Models [C]//2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS). Piscataway, NJ: IEEE, 2022: 996-1006. [12]WANG L N, YE J M, ZHAO Y Y, et al. Superneurons: Dynamic GPU Memory Management for Training Deep Neural Networks [C]//Proceedings of the 23rd ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. New York: ACM, 2018: 41-53. [13]PENG X, SHI X H, DAI H L, et al. Capuchin: Tensor-Based GPU Memory Management for Deep Learning [C]//Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems. New York: ACM, 2020: 891-905. (責(zé)任編輯:? 韓 嘯) 收稿日期: 2023-04-06. 第一作者簡介: 魏曉輝(1972—), 男, 漢族, 博士, 教授, 博士生導(dǎo)師, 從事分布式系統(tǒng)、 網(wǎng)格系統(tǒng)與網(wǎng)絡(luò)安全的研究, E-mail: weixh@jlu.edu.cn. 通信作者簡介: 李洪亮(1983—), 男, 漢族, 博士, 副教授, 博士生導(dǎo)師, 從事分布式系統(tǒng)與虛擬化的研究, E-mail: lihongliang@jlu.edu.cn. 基金項(xiàng)目: 吉林省自然科學(xué)基金面上項(xiàng)目(批準(zhǔn)號: 20230101062JC).