深度神經網絡模型訓練時GPU顯存不足怎麼辦？

前言

最近跑的模型都比較大，尤其是Bert， 這真的是難爲我 1080ti 了， 在Bert的Example中，官方提供了一些 Trick 來幫助我們加速訓練，很良心， 但感覺還不夠，於是花費一些時間整理出一個 Trick 集合，來幫助我們在顯存不足的時候來嘿嘿嘿。

本文分爲兩大部分，第一部分引入一個主題：如何估計模型所需顯存， 第二個主題：GPU顯存不足時的各種 Trick 。

監控 GPU

監控GPU最常用的當然是 nvidia-smi ，但有一個工具能夠更好的展示信息：gpustat 。

nvidia-smi
watch --color -n1 gpustat -cpu   # 動態事實監控GPU

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

下面有同學推薦nvtop， 我簡單試了試，的確挺好的，展現出現的信息很豐富 ， 推薦試一試。

如何估計模型顯存 [1]
首先，思考一個問題： 模型中的哪些東西佔據了我的顯存，咋就動不動就 out of memory？

其實一個模型所佔用的顯存主要包含兩部分： 模型自身的參數， 優化器參數， 模型每層的輸入輸出。

模型自身參數

模型自身的參數指的就是各個網絡層的 Weight 和Bias，這部分顯存在模型加載完成之後就會被佔用， 注意到的是，有些層是有參數的，如CNN， RNN； 而有些層是無參數的， 如激活層， 池化層等。

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

對於Pytorch來說，模型參數存儲在 model.parameters() 中，因此，我們不需要自己計算，完全可以通過Pytorh來直接打印：

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

優化器參數

優化器參數指的是模型在優化過程即反向傳播中所產生的參數， 這部分參數主要指的就是 dw， 即梯度，在SGD中， 其大小與參數一樣， 因此在優化期間， 模型的參數所佔用的顯存會翻倍。

值得注意的是，不同的優化器其所需保存的優化參數不同， 對於 Adam， 由於其還需要保存其餘參數， 模型的參數量會在優化區間翻 4 倍。

模型每層的輸入輸出

首先，第一點是輸入數據所佔用的顯存， 這部分所佔用的顯存其實並不大，這是因爲我們往往採用迭代器的方式讀取數據，這意味着我們其實並不是一次性的將所有數據讀入顯存，而這保證每次輸入所佔用的顯存與整個網絡參數來比是微不足道的。

然後，在模型進行前向傳播與反向傳播時， 一個很重要的事情就是計算並保存每一層的輸出以及其對應的梯度， 這意味着，這也佔據了很大一部分顯存。

最後， 模型輸出的顯存佔用可以總結爲：

• 每一層的輸出(多維數組)， 其對應的梯度， 值得注意的是，模型輸出不需要存儲相應的動量信息（即此處如果使用Adam， 模型輸出的參數量依舊是2倍而不是4倍， 我也不知道爲啥？？求大佬指教）
• 輸出的顯存佔用與 batch size 成正比

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

所有的顯存佔用計算

顯存佔用 = 模型自身參數 × n + batch size × 輸出參數量 × 2 + 一個batch的輸入數據（往往忽略）

其中，n是根據優化算法來定的，如果選用SGD， 則 n = 2， 如果選擇Adam， 則 n = 4.

一個很棒的實現如下， 我懶得再重新寫了，你可以根據這個改一改，問題不大。

# 模型顯存佔用監測函數
# model：輸入的模型
# input：實際中需要輸入的Tensor變量
# type_size 默認爲 4 默認類型爲 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 顯存不足時的Trick [2]

此處不討論多GPU， 分佈式計算等情況，只討論一些常規的 Trick， 會不定時進行更新。

降低batch size

這應該很好理解，適當降低batch size， 則模型每層的輸入輸出就會成線性減少， 效果相當明顯。這裏需要注意的一點是, dev batch size 的調整也有助於降低顯存, 同時,不要將 dev 或 test 的batch size 設置爲樣本集長度, 我最近就幹了這個傻事，害的我調試了一天才調出來是這個問題。

選擇更小的數據類型

一般默認情況下， 整個網絡中採用的是32位的浮點數，如果切換到 16位的浮點數，其顯存佔用量將接近呈倍數遞減。

精簡模型

在設計模型時，適當的精簡模型，如原來兩層的LSTM轉爲一層； 原來使用LSTM， 現在使用GRU； 減少卷積核數量； 儘量少的使用 Linear 等。

數據角度

對於文本數據來說，長序列所帶來的參數量是呈線性增加的， 適當的縮小序列長度可以極大的降低參數量。

total_loss

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

total_loss += loss.item()

釋放不需要的張量和變量

採用del釋放你不再需要的張量和變量，這也要求我們在寫模型的時候注意變量的使用，不要隨心所欲，漫天飛舞。

Relu 的 inplace 參數

激活函數 Relu() 有一個默認參數 inplace ，默認爲Flase， 當設置爲True的時候，我們在通過relu() 計算得到的新值不會佔用新的空間而是直接覆蓋原來的值，這表示設爲True， 可以節省一部分顯存。

梯度累積

首先， 要了解一些Pytorch的基本知識：

• 在Pytorch 中，當我們執行 loss.backward() 時， 會爲每個參數計算梯度，並將其存儲在 paramter.grad 中， 注意到， paramter.grad 是一個張量， 其會累加每次計算得到的梯度。
• 在 Pytorch 中， 只有調用 optimizer.step()時纔會進行梯度下降更新網絡參數。

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

我們先來看看傳統訓練：

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

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

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

比較來看， 我們發現，梯度累加本質上就是累加 accumulation_steps 個batch 的梯度， 再根據累加的梯度來更新網絡參數，以達到類似batch_size 爲 accumulation_steps * batch_size 的效果。在使用時，需要注意適當的擴大學習率。

更詳細來說， 我們假設 batch size = 32， accumulation steps = 8 ， 梯度積累首先在前向傳播的時候講 batch 分爲 accumulation steps 份， 然後得到 size=4 的小份batch ， 每次就以小 batch 來計算梯度，但是不更新參數，將梯度積累下來，直到我們計算了 accumulation steps 個小 batch， 我們再更新參數。

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

在Bert的倉庫中，就使用了這個Trick，十分實用，簡直是我們這種乞丐實驗室的良心Trick。

梯度檢查點

這個Trick我沒用過，畢竟模型還沒有那麼那麼大。

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

最後

哎， 如果你看完了這篇文章，就說明了一件事情： **小夥子，你卡也不夠啊。**哎， 乞丐實驗室不配深度學習，哭了。

Reference

[1]科普帖：深度學習中GPU和顯存分析

[2]如何在Pytorch中精細化利用顯存

[3]GPU捉襟見肘還想訓練大批量模型？誰說不可以

[4]PyTorch中在反向傳播前爲什麼要手動將梯度清零？

[5]From zero to research — An introduction to Meta-learning

[6]Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups