GPU顯存不足時(shí)的各種Trick

前言

最近跑的模型都比較大，尤其是Bert， 這真的是難為我 1080ti 了， 在Bert的Example中，官方提供了一些 Trick 來幫助我們加速訓(xùn)練，很良心， 但感覺還不夠，于是花費(fèi)一些時(shí)間整理出一個(gè) Trick 集合，來幫助我們?cè)陲@存不足的時(shí)候來嘿嘿嘿。

本文分為兩大部分，第一部分引入一個(gè)主題：如何估計(jì)模型所需顯存， 第二個(gè)主題：GPU顯存不足時(shí)的各種 Trick 。

監(jiān)控 GPU

監(jiān)控GPU最常用的當(dāng)然是nvidia-smi，但有一個(gè)工具能夠更好的展示信息：gpustat。

nvidia-smi watch --color -n1 gpustat -cpu # 動(dòng)態(tài)事實(shí)監(jiān)控GPU

推薦在配置文件中配置別名，反正我每次gpu一下，信息就全出來了，很方便。

下面有同學(xué)推薦nvtop， 我簡(jiǎn)單試了試，的確挺好的，展現(xiàn)出現(xiàn)的信息很豐富 ， 推薦試一試。

如何估計(jì)模型顯存 [1]

首先，思考一個(gè)問題：模型中的哪些東西占據(jù)了我的顯存，咋就動(dòng)不動(dòng)就out of memory？

其實(shí)一個(gè)模型所占用的顯存主要包含兩部分：模型自身的參數(shù)， 優(yōu)化器參數(shù)， 模型每層的輸入輸出。

模型自身參數(shù)

模型自身的參數(shù)指的就是各個(gè)網(wǎng)絡(luò)層的 Weight 和Bias，這部分顯存在模型加載完成之后就會(huì)被占用， 注意到的是，有些層是有參數(shù)的，如CNN， RNN；而有些層是無參數(shù)的， 如激活層， 池化層等。

從Pytorch 的角度來說，當(dāng)你執(zhí)行model.to(device)是， 你的模型就加載完畢，此時(shí)你的模型就已經(jīng)加載完成了。

對(duì)于Pytorch來說，模型參數(shù)存儲(chǔ)在model.parameters()中，因此，我們不需要自己計(jì)算，完全可以通過Pytorh來直接打?。?/p>

print('Model {} : params: {:4f}M'.format(model._get_name()， para * type_size / 1000 / 1000))

優(yōu)化器參數(shù)

優(yōu)化器參數(shù)指的是模型在優(yōu)化過程即反向傳播中所產(chǎn)生的參數(shù)， 這部分參數(shù)主要指的就是 dw， 即梯度，在SGD中， 其大小與參數(shù)一樣， 因此在優(yōu)化期間， 模型的參數(shù)所占用的顯存會(huì)翻倍。

值得注意的是，不同的優(yōu)化器其所需保存的優(yōu)化參數(shù)不同， 對(duì)于 Adam， 由于其還需要保存其余參數(shù)， 模型的參數(shù)量會(huì)在優(yōu)化區(qū)間翻 4 倍。

模型每層的輸入輸出

首先，第一點(diǎn)是輸入數(shù)據(jù)所占用的顯存， 這部分所占用的顯存其實(shí)并不大，這是因?yàn)槲覀兺捎玫鞯姆绞阶x取數(shù)據(jù)，這意味著我們其實(shí)并不是一次性的將所有數(shù)據(jù)讀入顯存，而這保證每次輸入所占用的顯存與整個(gè)網(wǎng)絡(luò)參數(shù)來比是微不足道的。

然后，在模型進(jìn)行前向傳播與反向傳播時(shí)， 一個(gè)很重要的事情就是計(jì)算并保存每一層的輸出以及其對(duì)應(yīng)的梯度， 這意味著，這也占據(jù)了很大一部分顯存。

最后，模型輸出的顯存占用可以總結(jié)為：

每一層的輸出(多維數(shù)組)， 其對(duì)應(yīng)的梯度， 值得注意的是，模型輸出不需要存儲(chǔ)相應(yīng)的動(dòng)量信息（即此處如果使用Adam， 模型輸出的參數(shù)量依舊是2倍而不是4倍， 我也不知道為啥？？求大佬指教）

輸出的顯存占用與 batch size 成正比

那么有沒有辦法通過Pytorch來計(jì)算這部分參數(shù)量呢？答案是有的，我們可以假設(shè)一個(gè)batch的樣本，然后通過model.modules()來對(duì)每一層進(jìn)行遍歷，獲得每一層的輸出shape， 然后就能夠獲得一個(gè)batch的數(shù)據(jù)的輸出參數(shù)量。[2]

所有的顯存占用計(jì)算

顯存占用 = 模型自身參數(shù) × n + batch size × 輸出參數(shù)量 × 2 + 一個(gè)batch的輸入數(shù)據(jù)（往往忽略）

其中，n是根據(jù)優(yōu)化算法來定的，如果選用SGD， 則 n = 2， 如果選擇Adam， 則 n = 4.

一個(gè)很棒的實(shí)現(xiàn)如下， 我懶得再重新寫了，你可以根據(jù)這個(gè)改一改，問題不大。

# 模型顯存占用監(jiān)測(cè)函數(shù) # model：輸入的模型 # input：實(shí)際中需要輸入的Tensor變量 # type_size 默認(rèn)為 4 默認(rèn)類型為 float32 def modelsize(model, input, type_size=4): para = sum([np.prod(list(p.size())) for p in model.parameters()]) print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000)) input_ = input.clone() input_.requires_grad_(requires_grad=False) mods = list(model.modules()) out_sizes = [] for i in range(1, len(mods)): m = mods[i] if isinstance(m, nn.ReLU): if m.inplace: continue out = m(input_) out_sizes.append(np.array(out.size())) input_ = out total_nums = 0 for i in range(len(out_sizes)): s = out_sizes[i] nums = np.prod(np.array(s)) total_nums += nums print('Model {} : intermedite variables: {:3f} M (without backward)' .format(model._get_name(), total_nums * type_size / 1000 / 1000)) print('Model {} : intermedite variables: {:3f} M (with backward)' .format(model._get_name(), total_nums * type_size*2 / 1000 / 1000))

GPU 顯存不足時(shí)的Trick [2]

此處不討論多GPU， 分布式計(jì)算等情況，只討論一些常規(guī)的 Trick， 會(huì)不定時(shí)進(jìn)行更新。

降低batch size

這應(yīng)該很好理解，適當(dāng)降低batch size， 則模型每層的輸入輸出就會(huì)成線性減少， 效果相當(dāng)明顯。這里需要注意的一點(diǎn)是, dev batch size的調(diào)整也有助于降低顯存, 同時(shí),不要將 dev 或 test 的batch size 設(shè)置為樣本集長(zhǎng)度, 我最近就干了這個(gè)傻事，害的我調(diào)試了一天才調(diào)出來是這個(gè)問題。

選擇更小的數(shù)據(jù)類型

一般默認(rèn)情況下， 整個(gè)網(wǎng)絡(luò)中采用的是32位的浮點(diǎn)數(shù)，如果切換到 16位的浮點(diǎn)數(shù)，其顯存占用量將接近呈倍數(shù)遞減。

精簡(jiǎn)模型

在設(shè)計(jì)模型時(shí)，適當(dāng)?shù)木?jiǎn)模型，如原來兩層的LSTM轉(zhuǎn)為一層；原來使用LSTM， 現(xiàn)在使用GRU；減少卷積核數(shù)量；盡量少的使用 Linear 等。

數(shù)據(jù)角度

對(duì)于文本數(shù)據(jù)來說，長(zhǎng)序列所帶來的參數(shù)量是呈線性增加的， 適當(dāng)?shù)目s小序列長(zhǎng)度可以極大的降低參數(shù)量。

total_loss

考慮到 loss 本身是一個(gè)包含梯度信息的 tensor， 因此，正確的求損失和的方式為：

total_loss += loss.item()

釋放不需要的張量和變量

采用del釋放你不再需要的張量和變量，這也要求我們?cè)趯懩Ｐ偷臅r(shí)候注意變量的使用，不要隨心所欲，漫天飛舞。

Relu 的 inplace 參數(shù)

激活函數(shù)Relu()有一個(gè)默認(rèn)參數(shù)inplace，默認(rèn)為Flase， 當(dāng)設(shè)置為True的時(shí)候，我們?cè)谕ㄟ^relu()計(jì)算得到的新值不會(huì)占用新的空間而是直接覆蓋原來的值，這表示設(shè)為True， 可以節(jié)省一部分顯存。

梯度累積

首先， 要了解一些Pytorch的基本知識(shí)：

在Pytorch 中，當(dāng)我們執(zhí)行l(wèi)oss.backward()時(shí)， 會(huì)為每個(gè)參數(shù)計(jì)算梯度，并將其存儲(chǔ)在 paramter.grad 中， 注意到，paramter.grad是一個(gè)張量， 其會(huì)累加每次計(jì)算得到的梯度。

在 Pytorch 中， 只有調(diào)用optimizer.step()時(shí)才會(huì)進(jìn)行梯度下降更新網(wǎng)絡(luò)參數(shù)。

我們知道， batch size 與占用顯存息息相關(guān)，但有時(shí)候我們的batch size 又不能設(shè)置的太小，這咋辦呢？答案就是梯度累加。

我們先來看看傳統(tǒng)訓(xùn)練：

for i,(feature,target) in enumerate(train_loader): outputs = model(feature) # 前向傳播 loss = criterion(outputs,target) # 計(jì)算損失 optimizer.zero_grad() # 清空梯度 loss.backward() # 計(jì)算梯度 optimizer.step() # 反向傳播， 更新網(wǎng)絡(luò)參數(shù)

而加入梯度累加之后，代碼是這樣的：

for i,(features,target) in enumerate(train_loader): outputs = model(images) # 前向傳播 loss = criterion(outputs,target) # 計(jì)算損失 loss = loss/accumulation_steps # 可選，如果損失要在訓(xùn)練樣本上取平均 loss.backward() # 計(jì)算梯度 if((i+1)%accumulation_steps)==0: optimizer.step() # 反向傳播，更新網(wǎng)絡(luò)參數(shù) optimizer.zero_grad() # 清空梯度

其實(shí)，這塊有兩種理解方式（受到評(píng)論區(qū)同學(xué)啟發(fā)）， 我談?wù)勗?bert 里面最常見的那種。

比較來看， 我們發(fā)現(xiàn)，梯度累加本質(zhì)上就是累加accumulation_steps個(gè)batchsize/accumulationsteps的梯度， 再根據(jù)累加的梯度來更新網(wǎng)絡(luò)參數(shù)，以達(dá)到真實(shí)梯度類似batch_size的效果。在使用時(shí)，需要注意適當(dāng)?shù)臄U(kuò)大學(xué)習(xí)率。

更詳細(xì)來說， 我們假設(shè)batch size = 4，accumulation steps = 8， 梯度積累首先在前向傳播的時(shí)候以batch_size=4來計(jì)算梯度，但是不更新參數(shù)，將梯度積累下來，直到我們計(jì)算了accumulation steps個(gè) batch， 我們?cè)俑聟?shù)。其實(shí)本質(zhì)上就等價(jià)于：

真正的 batch_size = batch_size * accumulation_steps

梯度積累能很大程度上緩解GPU顯存不足的問題，推薦使用。

在Bert的倉(cāng)庫(kù)中，就使用了這個(gè)Trick，十分實(shí)用，簡(jiǎn)直是我們這種乞丐實(shí)驗(yàn)室的良心Trick。

梯度檢查點(diǎn)

這個(gè)Trick我沒用過，畢竟模型還沒有那么那么大。

等我用過再更新吧，先把坑挖下。

最后

哎， 如果你看完了這篇文章，就說明了一件事情：小伙子，你卡也不夠啊。哎， 乞丐實(shí)驗(yàn)室不配深度學(xué)習(xí)，哭了。