結(jié)合算子選擇的卷積神經(jīng)網(wǎng)絡(luò)顯存優(yōu)化算法

魏曉輝　周博文　李洪亮　徐哲文

摘要： 針對(duì)卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練中自動(dòng)算子選擇算法在較大的顯存壓力下性能下降的問(wèn)題， 將卸載、 重計(jì)算與卷積算子選擇統(tǒng)一建模， 提出一種智能算子選擇算法。該算法權(quán)衡卸載和重計(jì)算引入的時(shí)間開(kāi)銷(xiāo)與更快的卷積算子節(jié)省的時(shí)間， 尋找卸載、 重計(jì)算和卷積算子選擇的調(diào)度， 解決了自動(dòng)算子選擇算法性能下降的問(wèn)題. 實(shí)驗(yàn)結(jié)果表明， 該智能算子選擇算法比重計(jì)算-自動(dòng)算子選擇算法縮短了13.53%訓(xùn)練時(shí)間， 比已有的卸載/重計(jì)算-自動(dòng)算子選擇算法縮短了4.36%的訓(xùn)練時(shí)間.

關(guān)鍵詞： 顯存； 卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練； 卷積算子； 卸載； 重計(jì)算

中圖分類(lèi)號(hào)： TP391文獻(xiàn)標(biāo)志碼： A文章編號(hào)： 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ò)卓越的性能， 越來(lái)越多的應(yīng)用使用了卷積神經(jīng)網(wǎng)絡(luò)， 如圖像識(shí)別、 語(yǔ)義分割等. 為加速卷積計(jì)算， 研究人員提出了不同的卷積算子， 它們需要不同的工作空間和卷積耗時(shí). 在訓(xùn)練時(shí)， 對(duì)于卷積算子的選擇， 主流的訓(xùn)練框架如Tensorflow［1］， 會(huì)根據(jù)空閑顯存自動(dòng)選擇合適的卷積算子. 但卷積神經(jīng)網(wǎng)絡(luò)的訓(xùn)練需要保存大量的張量以計(jì)算梯度， 訓(xùn)練的顯存壓力隨網(wǎng)絡(luò)規(guī)模的增大而增加. 當(dāng)顯存壓力變大， 空閑顯存匱乏時(shí)， 已有的自動(dòng)算子選擇算法就會(huì)選擇工作空間需求少但耗時(shí)長(zhǎng)的卷積算子， 進(jìn)而延長(zhǎng)了訓(xùn)練時(shí)間.

目前緩解訓(xùn)練顯存壓力的主流方法是卸載和重計(jì)算張量. 卸載方法［2-5］將顯存中暫時(shí)不用的張量通過(guò)PCIE總線傳輸?shù)饺萘扛蟮腃PU內(nèi)存， 等再次使用時(shí)再?gòu)腃PU的內(nèi)存?zhèn)骰谿PU的顯存， 從而緩解訓(xùn)練的存儲(chǔ)壓力， 運(yùn)行效率受PCIE總線傳輸速度的影響. 重計(jì)算方法［6-7］保留部分張量作為檢查點(diǎn)， 等使用被釋放的張量時(shí)再?gòu)臋z查點(diǎn)重新計(jì)算回來(lái)， 運(yùn)行效率受檢查點(diǎn)機(jī)制的影響. 通過(guò)卸載和重計(jì)算釋放更多的顯存， 本文選擇工作空間大但耗時(shí)短的卷積算子， 可有效減少訓(xùn)練時(shí)間， 但也會(huì)引入更多的卸載和重計(jì)算時(shí)間開(kāi)銷(xiāo). 所以將卸載、 重計(jì)算和算子選擇開(kāi)銷(xiāo)統(tǒng)一建模， 并基于開(kāi)銷(xiāo)模型進(jìn)行綜合考慮設(shè)計(jì)算法已經(jīng)成為目前該領(lǐng)域重要的研究?jī)?nèi)容.

本文針對(duì)卷積算子選擇的問(wèn)題， 進(jìn)行整數(shù)規(guī)劃問(wèn)題的模型設(shè)計(jì). 在此基礎(chǔ)上， 本文提出智能算子選擇算法， 綜合考慮卸載、 重計(jì)算和卷積算子選擇， 使卸載和重計(jì)算為更快的卷積算子的工作空間騰出顯存， 進(jìn)而縮短訓(xùn)練時(shí)間. 本文使用大規(guī)模求解器Gurobi搜索整數(shù)規(guī)劃模型的結(jié)果作為建模整數(shù)規(guī)劃問(wèn)題的可行解， 并將求解器結(jié)果轉(zhuǎn)化為顯存管理調(diào)度， 在訓(xùn)練時(shí)通過(guò)調(diào)度管理顯存. 本文對(duì)實(shí)際AI生產(chǎn)環(huán)境下的多種網(wǎng)絡(luò)進(jìn)行測(cè)試， 并將智能算子選擇與自動(dòng)算子選擇進(jìn)行對(duì)比. 實(shí)驗(yàn)結(jié)果表明， 本文的MO-SOS（memory optimization with smart operator selection）算法可以顯著提升GPU運(yùn)算性能.

1 研究背景

1.1 自動(dòng)算子選擇與卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練顯存壓力

由于計(jì)算量大， 卷積計(jì)算是卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練中最耗時(shí)的計(jì)算. 為加速卷積計(jì)算， 研究人員提出了不同的卷積算子實(shí)現(xiàn)卷積計(jì)算， 如快速Fourier變換卷積［8］和Winograd變換卷積［9］等. 它們?cè)诩铀倬矸e計(jì)算時(shí)需額外的空間保存計(jì)算過(guò)程中的臨時(shí)變量， 即神經(jīng)網(wǎng)絡(luò)算子的工作空間. 在訓(xùn)練卷積神經(jīng)網(wǎng)絡(luò)時(shí)， 現(xiàn)有的訓(xùn)練框架如Tensorflow， 會(huì)根據(jù)空閑顯存自動(dòng)選擇合適的卷積算子.

近年來(lái)， 卷積神經(jīng)網(wǎng)絡(luò)的網(wǎng)絡(luò)規(guī)模越來(lái)越大， 使訓(xùn)練的顯存壓力變得越來(lái)越嚴(yán)峻. 訓(xùn)練卷積神經(jīng)網(wǎng)絡(luò)是尋找一組網(wǎng)絡(luò)參數(shù)使神經(jīng)網(wǎng)絡(luò)的輸出盡可能接近期望值. 圖1為一個(gè)多層卷積神經(jīng)網(wǎng)絡(luò)的訓(xùn)練過(guò)程， 包括卷積層CONV、 批標(biāo)準(zhǔn)化層BN、 激活層ReLU和全連接層FC等. 目前尋找網(wǎng)絡(luò)參數(shù)常用的方法是梯度下降法， 需要計(jì)算神經(jīng)網(wǎng)絡(luò)的損失（Loss）和參數(shù)的梯度.

Loss是神經(jīng)網(wǎng)絡(luò)輸出與期望值之間的差距， 這種差距可通過(guò)計(jì)算損失函數(shù)得到， 即圖1中的Loss層. Loss是通過(guò)前向傳播計(jì)算的， 前向傳播是圖1中實(shí)線表示的方向. 神經(jīng)網(wǎng)絡(luò)訓(xùn)練可通過(guò)Loss計(jì)算出模型的參數(shù)梯度， 從而使用參數(shù)梯度更新模型參數(shù).

計(jì)算梯度是通過(guò)后向傳播完成的， 即圖1中虛線表示的方向. 計(jì)算梯度使用鏈?zhǔn)椒▌t依次計(jì)算出每層的參數(shù)梯度. 計(jì)算參數(shù)梯度依賴(lài)張量[WTHX]a[WTBX]i， 所以卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練會(huì)保存前向傳播的輸出張量. 目前， 卷積神經(jīng)網(wǎng)絡(luò)通過(guò)增加層數(shù)以獲得更高準(zhǔn)確率， 所以它們的訓(xùn)練顯存需求也會(huì)越來(lái)越大， 如Vgg19［10］和Resnet1001等. 當(dāng)顯存壓力較大， 空閑顯存匱乏時(shí)， 自動(dòng)算子選擇算法會(huì)選擇工作空間少但耗時(shí)較長(zhǎng)的卷積算子， 增加訓(xùn)練時(shí)間. 但卷積神經(jīng)網(wǎng)絡(luò)的張量在訓(xùn)練中大部分時(shí)間內(nèi)都是空閑的， 未被訪問(wèn)［2］， 所以可將張量卸載或重計(jì)算減輕訓(xùn)練顯存壓力.

1.2 緩解顯存壓力方案與智能卷積算子選擇動(dòng)機(jī)

卸載將CPU的內(nèi)存作為GPU的緩存， 暫存訓(xùn)練中的張量. 目前的GPU通常作為CPU的外設(shè)設(shè)備， 它是通過(guò)PCIE總線與CPU相連， 它們通過(guò)PCIE總線相互傳輸數(shù)據(jù). 由于CPU的內(nèi)存比GPU的顯存容量大， 所以可以把張量卸載到CPU主存， 再次需要張量時(shí)預(yù)取回GPU顯存. 目前的主流顯卡如3090，P40等， 都支持計(jì)算、 卸載和預(yù)取張量并行. 所以為不阻塞訓(xùn)練中的計(jì)算， 需要提前卸載或預(yù)取張量. VDNN（virtualized deep neural networks）［2］在前向傳播中的卷積計(jì)算時(shí)卸載它的輸入張量， 當(dāng)后向傳播需要已卸載的輸入張量時(shí)再提前預(yù)取回. SwapAdvisor［4］通過(guò)進(jìn)化計(jì)算選擇卸載的張量、 卸載開(kāi)始時(shí)間和預(yù)取開(kāi)始時(shí)間， 進(jìn)一步減少阻塞時(shí)間. STR（swap dominated tensor re-generation）［11］發(fā)現(xiàn)以前的工作未考慮到多次預(yù)取張量， 因?yàn)榘岩粋€(gè)張量卸載到CPU的內(nèi)存后， CPU的內(nèi)存中便有了該張量的備份， 當(dāng)在訓(xùn)練中多次訪問(wèn)這個(gè)張量時(shí)， STR認(rèn)為可以把張量多次預(yù)取回GPU顯存.

隨著GPU計(jì)算能力越來(lái)越強(qiáng)， 張量的卸載和預(yù)取的阻塞可能比計(jì)算更耗時(shí)， 因此研究人員又提出了另一種緩解顯存壓力的方法——重計(jì)算. 重計(jì)算保留一部分張量作為檢查點(diǎn)， 釋放其他的張量騰出顯存空間， 當(dāng)被釋放的張量再次需要訪問(wèn)時(shí)， 從現(xiàn)有的檢查點(diǎn)開(kāi)始重新計(jì)算. 由于引入了重計(jì)算， 訓(xùn)練會(huì)引入額外的計(jì)算時(shí)間， 該方法是一種用時(shí)間換取空間的方法， 所以保留不同的檢查點(diǎn)會(huì)引入不同的重計(jì)算時(shí)間. Chen等［6］把網(wǎng)絡(luò)分成多個(gè)部分， 每部分只保留其輸入， 對(duì)于每部分的其他張量都重新計(jì)算. Checkmate［7］不同于文獻(xiàn)［6］的工作， 其探究了在顯存容量的限制下， 哪些張量作為檢查點(diǎn)能引入最小的額外計(jì)算時(shí)間， 將重計(jì)算問(wèn)題抽象為線性規(guī)劃問(wèn)題， 用現(xiàn)有的求解器解決重計(jì)算問(wèn)題.

卸載和預(yù)取會(huì)導(dǎo)致額外的阻塞時(shí)間， 而重計(jì)算會(huì)引入額外的計(jì)算時(shí)間. 文獻(xiàn)［12］認(rèn)為卷積計(jì)算時(shí)間較長(zhǎng)， 不適合用于重計(jì)算， 所以卸載和預(yù)取卷積計(jì)算的輸出張量， 而對(duì)于激活、 批標(biāo)準(zhǔn)化的輸出張量， 它們的計(jì)算時(shí)間相對(duì)較少， 適合用于重計(jì)算. 文獻(xiàn)［13］研究表明， 以往的卸載和重計(jì)算的結(jié)合都是以神經(jīng)網(wǎng)絡(luò)層的粒度管理顯存， 然而一個(gè)網(wǎng)絡(luò)層中也會(huì)存在多個(gè)計(jì)算， 所以提出了張量粒度的顯存管理方案. STR認(rèn)為在非線性卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練時(shí)， 一個(gè)張量會(huì)被多次訪問(wèn)， 所以STR通過(guò)多次預(yù)取已卸載的張量減少重計(jì)算的時(shí)間.

當(dāng)顯存壓力較大， 空閑顯存匱乏時(shí)， 可通過(guò)卸載和重計(jì)算釋放更多的顯存， 為卷積算子增加空閑顯存. 如圖2所示， 它們是Vgg19在Tesla P40平臺(tái)上訓(xùn)練時(shí)的6個(gè)前向傳播的卷積計(jì)算， 使用了主流的CUDNN作為卷積算子庫(kù). 不可釋放顯存是當(dāng)前卷積計(jì)算依賴(lài)的輸入和輸出張量占用的顯存， 而可釋放顯存是暫時(shí)不用的張量顯存. 與現(xiàn)有的工作相同， 本文主要對(duì)前向傳播的輸出張量進(jìn)行顯存優(yōu)化， 所以圖2只統(tǒng)計(jì)了前向傳播的輸出張量. 在訓(xùn)練時(shí)， 卸載和重計(jì)算可釋放的顯存， 為卷積計(jì)算騰出空閑顯存. 但卸載和重計(jì)算會(huì)引入更多的額外時(shí)間開(kāi)銷(xiāo)， 所以將卸載和重計(jì)算引入的時(shí)間開(kāi)銷(xiāo)與更快卷積算子節(jié)省的時(shí)間權(quán)衡. 本文拓展了STR的模型， 對(duì)卸載、 重計(jì)算和卷積算子選擇進(jìn)行統(tǒng)一建模， 尋找卸載、 重計(jì)算和卷積算子的調(diào)度， 以減少訓(xùn)練時(shí)間.

2 問(wèn)題模型與算法設(shè)計(jì)

2.1 考慮卷積算子選擇的顯存優(yōu)化問(wèn)題描述

由于深度神經(jīng)網(wǎng)絡(luò)框架如Tensorflow和Pytorch等的使用方便與性能強(qiáng)大， 神經(jīng)網(wǎng)絡(luò)設(shè)計(jì)者常使用它們?cè)O(shè)計(jì)和訓(xùn)練網(wǎng)絡(luò). 這些主流的框架將深度神經(jīng)網(wǎng)絡(luò)的訓(xùn)練過(guò)程表示成計(jì)算圖， 如有向無(wú)環(huán)圖（DAG）. 計(jì)算圖由節(jié)點(diǎn)和邊組成， 節(jié)點(diǎn)表示計(jì)算， 如卷積、 矩陣相乘等， 邊表示數(shù)據(jù)依賴(lài)， 如圖3中的DAG計(jì)算圖. 在訓(xùn)練時(shí)， 訓(xùn)練框架會(huì)按照計(jì)算圖的拓?fù)渑判颍?依次執(zhí)行節(jié)點(diǎn)的計(jì)算［4，7］.

圖3中v18節(jié)點(diǎn)是Vgg19模型的一個(gè)卷積計(jì)算， 運(yùn)行在Tesla P40平臺(tái)上， 算子庫(kù)使用了8.0.5版本的CUDNN. 當(dāng)顯存壓力變大， 空閑顯存匱乏時(shí)， 自動(dòng)算子選擇只能選擇不需要工作空間但耗時(shí)最長(zhǎng)的算子Implicit GEMM， 增加訓(xùn)練耗費(fèi)的時(shí)間. 所以MO-SOS算法要通過(guò)卸載和重計(jì)算為工作空間釋放顯存， 選擇更快的卷積算子. 選擇最快的算子Winograd Nonfused， 則需要通過(guò)卸載和重計(jì)算為工作空間釋放約2 GB的顯存. 相應(yīng)地， 卸載和重計(jì)算也會(huì)引入更多的額外時(shí)間開(kāi)銷(xiāo). 若引入的額外時(shí)間開(kāi)銷(xiāo)大于最快卷積算子節(jié)省的訓(xùn)練時(shí)間， 反而會(huì)增加訓(xùn)練時(shí)間. 另一種算子的選擇是工作空間少但耗時(shí)較長(zhǎng)的卷積算子Winograd. Winograd算子只需要通過(guò)卸載和重計(jì)算為工作空間釋放約6 MB的顯存.

選擇最快或最慢的卷積算子都有可能增加訓(xùn)練時(shí)間， 所以需要對(duì)卷積算子、 卸載和重計(jì)算進(jìn)行統(tǒng)一建模， 尋找一個(gè)卷積算子和卸載/重計(jì)算的調(diào)度， 使得從v1計(jì)算節(jié)點(diǎn)到vn計(jì)算節(jié)點(diǎn)執(zhí)行時(shí)間最短. 針對(duì)卷積算子的選擇問(wèn)題， 本文從考慮卷積算子選擇問(wèn)題的定義、 卸載和預(yù)取建模及工作空間的顯存限制三方面進(jìn)行問(wèn)題模型的設(shè)計(jì).

2.2 考慮卷積算子選擇問(wèn)題的定義

設(shè)一個(gè)訓(xùn)練計(jì)算圖為G=（V，E）， 其中V={v1，v2，…，vn}表示該計(jì)算圖有n個(gè)計(jì)算節(jié)點(diǎn)， E表示計(jì)算節(jié)點(diǎn)之間的依賴(lài)關(guān)系. 每個(gè)計(jì)算節(jié)點(diǎn)vi的輸出需要Mi顯存存儲(chǔ)輸出張量. 每個(gè)計(jì)算節(jié)點(diǎn)vi依賴(lài)于其他計(jì)算節(jié)點(diǎn)的輸出vj

/St+1，j≤∑aj/v=1Rt，j，v+Pss（t，j），j， t，j，（5）其中： ?ak表示k計(jì)算節(jié)點(diǎn)有多少種算子實(shí)現(xiàn); O和P變量表示卸載和預(yù)取的決策; 式（2）表示v1到vn計(jì)算節(jié)點(diǎn)都要計(jì)算一次; 式（3）約束了對(duì)于不同的計(jì)算節(jié)點(diǎn)， 算子的選擇最多只能有一種. 如上所述， 每個(gè)計(jì)算節(jié)點(diǎn)都依賴(lài)于其他計(jì)算節(jié)點(diǎn)的輸出張量. 本文使用卸載和重計(jì)算釋放顯存空間， 在某個(gè)節(jié)點(diǎn)計(jì)算時(shí)， 它依賴(lài)的張量可能并不存在于顯存中， 它可能未重計(jì)算或者預(yù)取回顯存. 所以MO-SOS算法使用St，j變量表示在t階段時(shí)， j張量是否可以訪問(wèn). 式（4）表示在計(jì)算每個(gè)節(jié)點(diǎn)時(shí)它依賴(lài)的張量必須可以訪問(wèn)， DEPS［k］表示k計(jì)算節(jié)點(diǎn)的依賴(lài)張量. 而式（5）表示如果在t階段重計(jì)算或者預(yù)取回了j張量， 則j張量在之后的計(jì)算中能被其他計(jì)算節(jié)點(diǎn)訪問(wèn).

2.3 卸載和預(yù)取建模

由于卸載、 預(yù)取與重計(jì)算可以并行， 因此MO-SOS算法可通過(guò)卸載和預(yù)取減少需要重計(jì)算的張量. 但為保證卸載與預(yù)取的正確性， MO-SOS算法需要對(duì)卸載和預(yù)取進(jìn)行約束. 本文規(guī)定在每個(gè)節(jié)點(diǎn)計(jì)算的開(kāi)端， 開(kāi)始卸載或預(yù)取張量. Ot，j表示在t階段是否開(kāi)始卸載j張量， Pt，j表示在t階段是否開(kāi)始預(yù)取j張量， 則∑n/t=1Ot，j≤1， j，（6）

式（6）表示每個(gè)張量只能卸載一次， 因?yàn)閺埩勘恍遁d到CPU內(nèi)存后， 內(nèi)存中的張量可作為備份多次預(yù)取. 式（7）表示在卸載時(shí)張量的數(shù)據(jù)必須是可以訪問(wèn)的， 否則會(huì)傳輸臟數(shù)據(jù). ss（t，i）表示在t階段完整地卸載或預(yù)取i張量， 需要在哪個(gè)階段開(kāi)始卸載或預(yù)取； sf（t，i）表示在t階段開(kāi)始卸載或預(yù)取i張量， 會(huì)在哪個(gè)階段完成. 這兩個(gè)函數(shù)可以根據(jù)PCIE總線帶寬、 張量大小和每個(gè)節(jié)點(diǎn)的計(jì)算時(shí)間計(jì)算得出. 但不同的卷積算子耗費(fèi)不同的卷積時(shí)間. 為解決上述問(wèn)題， 在計(jì)算ss和sf函數(shù)時(shí)， 可根據(jù)最快的卷積算子計(jì)算， 保證卸載和預(yù)取不會(huì)阻塞計(jì)算. 式（8）表示如果沒(méi)有預(yù)取就不會(huì)卸載張量， 且m是可設(shè)置的參數(shù). 式（9）表示預(yù)取必須在張量卸載完成后開(kāi)始. 此外， 本文限制了每個(gè)階段只能卸載或預(yù)取一個(gè)張量， 即Ot，j≤∑

2.4 考慮工作空間的顯存限制

為約束每個(gè)節(jié)點(diǎn)運(yùn)行時(shí)顯存不會(huì)溢出， 需要定義一個(gè)變量表示每個(gè)節(jié)點(diǎn)計(jì)算時(shí)需要的顯存. 所以用Ut，k表示在t階段計(jì)算k節(jié)點(diǎn)時(shí)需要的顯存， Ut，0表示在t階段開(kāi)始時(shí)需要的顯存， 它包括模型的輸入、 參數(shù)、 梯度以及作為檢查點(diǎn)的張量Ut，0=Minput+Mparam+Mgrad+∑n/j=1St，jMj+∑（x，y）∈OCC（t）MyPx，y.（14）此外， 預(yù)取也會(huì)占用顯存， 所以O(shè)CC（t）表示了哪些張量的預(yù)取在t階段占用顯存. 然后遞歸地計(jì)算t階段每個(gè)節(jié)點(diǎn)計(jì)算時(shí)需要的顯存Ut，k：Ut，k+1=Ut，k-mem_free（k）+∑ak+1/v=1Rt，k+1，vMk+1+∑

MO-SOS算法希望通過(guò)卸載和重計(jì)算騰出更多的顯存空間， 以選擇更快的卷積算子， 進(jìn)而減少訓(xùn)練時(shí)間. 式（2）～（5）保證了重計(jì)算的正確性， 式（6）～（13）保證了卸載和預(yù)取的正確性， 且式（14）～（18）約束了顯存占用. 該問(wèn)題是難以解決的NP-難問(wèn)題， 求解空間巨大， 無(wú)法通過(guò)高復(fù)雜度的遍歷算法求解， 因此需要有針對(duì)性地進(jìn)行模型設(shè)計(jì).

2.5 算法設(shè)計(jì)

本文問(wèn)題可以使用求解器根據(jù)輸入的計(jì)算圖和GPU的顯存限制求解一個(gè)可行的方案. 目標(biāo)公式（1）中的R，S，O，P， Free都是0～1變量， 所以是一個(gè)整數(shù)規(guī)劃問(wèn)題. MO-SOS算法使用現(xiàn)有的求解器去搜索一個(gè)可行的結(jié)果， 例如Gurobi，Cplex，Scip等. Gurobi是大規(guī)模數(shù)學(xué)規(guī)劃優(yōu)化器， 在多個(gè)整數(shù)規(guī)劃問(wèn)題中有更快的優(yōu)化速度， 適合于本文的最小化問(wèn)題.

根據(jù)Gurobi求解器給出的結(jié)果， MO-SOS算法會(huì)產(chǎn)生一個(gè)顯存管理調(diào)度L={l1，l2，…，lk}， 其中l(wèi)i是顯存管理的操作， 如計(jì)算張量、 釋放張量、 卸載張量和預(yù)取張量. MO-SOS算法如下.

MO-SOS算法遍歷所有的R，O，P， Free變量去制定顯存管理調(diào)度. 如果Ri，j，v為1， 則在i階段， MO-SOS算法要計(jì)算j張量， 添加顯存管理操作計(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ì)算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í)使用數(shù)據(jù)集Imagenet， 圖片大小為227×227， Unet在訓(xùn)練時(shí)使用數(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ī)劃問(wèn)題的求解器使用Gurobi， 版本為9.5.2.

本文將智能算子選擇算法MO-SOS和自動(dòng)算子選擇算法Auto進(jìn)行比較. 為能在顯存壓力較大的情況下進(jìn)行實(shí)驗(yàn)， 給自動(dòng)算子選擇算法配置兩種最新的顯存優(yōu)化方法： Checkmate（只有重計(jì)算）和STR（卸載和重計(jì)算相結(jié)合）. 這兩種顯存優(yōu)化算法與自動(dòng)算子選擇結(jié)合后分別為Checkmate-Auto和STR-Auto算法. 本文設(shè)置了不同的顯存比例， 逐漸減少能使用的顯存， 增加訓(xùn)練的顯存壓力， 測(cè)試了Checkmate-Auto，STR-Auto和MO-SOS算法的訓(xùn)練時(shí)間. 由于訓(xùn)練是多個(gè)訓(xùn)練步組成的， 每個(gè)訓(xùn)練步執(zhí)行時(shí)間浮動(dòng)較小， 所以本文對(duì)每個(gè)神經(jīng)網(wǎng)絡(luò)測(cè)試了20個(gè)訓(xùn)練步， 取平均值作為每個(gè)神經(jīng)網(wǎng)絡(luò)一個(gè)訓(xùn)練步的執(zhí)行時(shí)間. 本文測(cè)試了在不同的顯存壓力下， 智能算子選擇算法和自動(dòng)算子選擇算法訓(xùn)練步的執(zhí)行時(shí)間， 以考察智能算子選擇算法能否縮短訓(xùn)練時(shí)間.

[HT5H〗3.2 實(shí)驗(yàn)結(jié)果分析

圖4為不同算法對(duì)Vgg19和Unet一個(gè)訓(xùn)練步的執(zhí)行時(shí)間. 由圖4可見(jiàn)， 智能算子選擇算法在多種顯存壓力下都優(yōu)于自動(dòng)算子選擇算法. 本文的智能算子選擇算法MO-SOS平均比只有重計(jì)算的Checkmate-Auto算法減少了13.53%的訓(xùn)練時(shí)間， 比卸載和重計(jì)算結(jié)合的STR-Auto算法減少了4.36%的訓(xùn)練時(shí)間. 在顯存比例為0.8和0.7， 顯存壓力較小時(shí)， MO-SOS和STR-Auto算法的訓(xùn)練時(shí)間基本相當(dāng)， 而結(jié)合了卸載的MO-SOS和STR-Auto算法仍比Checkmate-Auto算法減少約11.2%的訓(xùn)練時(shí)間， 這主要是前者使用卸載減少了需要重計(jì)算的張量.

隨著顯存壓力變大， MO-SOS算法的訓(xùn)練時(shí)間逐漸比STR-Auto算法減少. 這是因?yàn)轱@存壓力變大， STR-Auto算法沒(méi)有為更快的卷積算子釋放更多的顯存， 選擇了更慢的卷積算子， 增加了訓(xùn)練時(shí)間. 但對(duì)于Unet網(wǎng)絡(luò)， 在顯存比例為0.4時(shí)， MO-SOS與STR-Auto算法的訓(xùn)練時(shí)間差距縮小. 不同于Vgg19， Unet是非線性網(wǎng)絡(luò)， 前向傳播的輸出張量不僅在后向傳播計(jì)算梯度時(shí)多次使用， 并且在前向傳播中也會(huì)多次使用. 所以Unet在訓(xùn)練時(shí)需要長(zhǎng)時(shí)間保存多個(gè)張量， 導(dǎo)致更大的顯存壓力， 難以選擇更快的卷積算子.

MO-SOS和STR-Auto算法訓(xùn)練時(shí)間比Checkmate-Auto算法少， 這是因?yàn)榍罢咄ㄟ^(guò)卸載減少了重計(jì)算的張量. 如圖5所示， 在Vgg19訓(xùn)練時(shí)， MO-SOS和STR-Auto算法在顯存比例為0.7時(shí)只有卸載， 而Checkmate-Auto算法只有重計(jì)算. MO-SOS和STR-Auto算法通過(guò)卸載騰出了足夠的顯存， 所以它們的卷積計(jì)算和其他計(jì)算的耗時(shí)都明顯少于Checkmate-Auto算法.

MO-SOS算法訓(xùn)練時(shí)間比STR-Auto算法少， 是因?yàn)镸O-SOS算法卸載和重計(jì)算更多的張量以增加空閑顯存， 選擇更快的卷積算子減少訓(xùn)練時(shí)間. 由圖5可見(jiàn)， 當(dāng)顯存比例為0.7時(shí)， MO-SOS算法比STR-Auto算法卸載的張量更多. 所以MO-SOS算法的卷積計(jì)算耗時(shí)少于STR-Auto算法. 隨著顯存壓力增大， MO-SOS和STR-Auto算法逐漸增加了重計(jì)算的張量以釋放顯存. 當(dāng)顯存比例為0.5時(shí)， 為給卷積計(jì)算的工作空間騰出更多的顯存， MO-SOS算法卸載和重計(jì)算的張量多于STR-Auto算法. 所以MO-SOS算法的其他計(jì)算耗時(shí)多于STR-Auto算法， 但MO-SOS算法的卷積計(jì)算耗時(shí)明顯少于STR-Auto算法， 總體訓(xùn)練時(shí)間也少于STR-Auto算法.

實(shí)驗(yàn)結(jié)果表明， 可以通過(guò)卸載和重計(jì)算為卷積計(jì)算增加空閑顯存， 選擇更快的卷積算子減少訓(xùn)練時(shí)間. 在顯存壓力較小時(shí)， 減少的訓(xùn)練時(shí)間并不明顯. 但隨著顯存壓力增大， 自動(dòng)算子選擇算法的空閑顯存逐漸減少， 選擇了更慢的卷積算子. 而智能算子選擇算法通過(guò)卸載和重計(jì)算釋放了更多的顯存， 選擇了更快的卷積算子， 進(jìn)而縮短了訓(xùn)練時(shí)間.

綜上所述， 針對(duì)自動(dòng)算子選擇算法在空閑顯存匱乏時(shí)會(huì)選擇工作空間需求少但耗時(shí)長(zhǎng)的卷積算子， 降低了訓(xùn)練的性能、 增加了訓(xùn)練時(shí)間的問(wèn)題， 本文提出了一種智能算子選擇算法， 通過(guò)卸載和重計(jì)算增加空閑顯存. 本文對(duì)卸載、 重計(jì)算和卷積算子選擇進(jìn)行了統(tǒng)一的整數(shù)規(guī)劃建模， 使用現(xiàn)有的求解器進(jìn)行求解， 并對(duì)多種網(wǎng)絡(luò)在各種顯存壓力下進(jìn)行實(shí)驗(yàn). 實(shí)驗(yàn)結(jié)果表明， 智能算子選擇算法在多種顯存壓力下優(yōu)于自動(dòng)算子選擇算法.

參考文獻(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.

第一作者簡(jiǎn)介： 魏曉輝（1972—）， 男， 漢族， 博士， 教授， 博士生導(dǎo)師， 從事分布式系統(tǒng)、 網(wǎng)格系統(tǒng)與網(wǎng)絡(luò)安全的研究， E-mail： weixh@jlu.edu.cn.

通信作者簡(jiǎn)介： 李洪亮（1983—）， 男， 漢族， 博士， 副教授， 博士生導(dǎo)師， 從事分布式系統(tǒng)與虛擬化的研究， E-mail： lihongliang@jlu.edu.cn.

基金項(xiàng)目： 吉林省自然科學(xué)基金面上項(xiàng)目（批準(zhǔn)號(hào)： 20230101062JC）.