淺談深度學習:如何計算模型以及中間變量的顯存佔用大小
前言
親，顯存炸了，你的顯卡快冒煙了！
torch.FatalError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/generic/THCStorage.cu:58
想必這是所有煉丹師們最不想看到的錯誤，沒有之一。
OUT OF MEMORY，顯然是顯存裝不下你那麼多的模型權重還有中間變量，然後程序奔潰了。怎麼辦，其實辦法有很多，及時清空中間變量，優化代碼，減少batch，等等等等，都能夠減少顯存溢出的風險。
但是這篇要說的是上面這一切優化操作的基礎，如何去計算我們所使用的顯存。學會如何計算出來我們設計的模型以及中間變量所佔顯存的大小，想必知道了這一點，我們對自己顯存也就會得心應手了。
如何計算
首先我們應該了解一下基本的數據量信息：
- 1 G = 1000 MB
- 1 M = 1000 KB
- 1 K = 1000 Byte
- 1 B = 8 bit
好，肯定有人會問為什麼是1000而不是1024，這裡不過多討論，只能說兩種說法都是正確的，只是應用場景略有不同。這裡統一按照上面的標準進行計算。
然後我們說一下我們平常使用的向量所佔的空間大小，以Pytorch官方的數據格式為例（所有的深度學習框架數據格式都遵循同一個標準）：
我們只需要看左邊的信息，在平常的訓練中，我們經常使用的一般是這兩種類型：
- float32 單精度浮點型
- int32 整型
一般一個8-bit的整型變量所佔的空間為1B也就是8bit。而32位的float則占4B也就是32bit。而雙精度浮點型double和長整型long在平常的訓練中我們一般不會使用。
ps：消費級顯卡對單精度計算有優化，服務器級別顯卡對雙精度計算有優化。
也就是說，假設有一幅RGB三通道真彩色圖片，長寬分別為500 x 500，數據類型為單精度浮點型，那麼這張圖所佔的顯存的大小為：500 x 500 x 3 x 4B = 3M。
而一個(256,3,100,100)-(N,C,H,W)的FloatTensor所佔的空間為256 x 3 x 100 x 100 x 4B = 31M
不多是吧，沒關係，好戲才剛剛開始。
顯存去哪兒了
看起來一張圖片(3x256x256)和卷積層(256x100x100)所佔的空間並不大，那為什麼我們的顯存依舊還是用的比較多，原因很簡單，佔用顯存比較多空間的並不是我們輸入圖像，而是神經網絡中的中間變量以及使用optimizer算法時產生的巨量的中間參數。
我們首先來簡單計算一下Vgg16這個net需要佔用的顯存：
通常一個模型佔用的顯存也就是兩部分：
- 模型自身的參數(params)
- 模型計算產生的中間變量(memory)
圖片來自cs231n，這是一個典型的sequential-net，自上而下很順暢，我們可以看到我們輸入的是一張224x224x3的三通道圖像，可以看到一張圖像只佔用150x4k，但上面是150k，這是因為這裡在計算的時候默認的數據格式是8-bit而不是32-bit，所以最後的結果要乘上一個4。
我們可以看到，左邊的memory值代表：圖像輸入進去，圖片以及所產生的中間卷積層所佔的空間。我們都知道，這些形形色色的深層卷積層也就是深度神經網絡進行「思考」的過程：
圖片從3通道變為64 –> 128 –> 256 –> 512 …. 這些都是卷積層，而我們的顯存也主要是他們佔用了。
還有上面右邊的params，這些是神經網絡的權重大小，可以看到第一層卷積是3×3，而輸入圖像的通道是3，輸出通道是64，所以很顯然，第一個卷積層權重所佔的空間是 (3 x 3 x 3) x 64。
另外還有一個需要注意的是中間變量在backward的時候會翻倍！
舉個例子，下面是一個計算圖，輸入x，經過中間結果z，然後得到最終變量L：
我們在backward的時候需要保存下來的中間值。輸出是L，然後輸入x，我們在backward的時候要求L對x的梯度，這個時候就需要在計算鏈L和x中間的z：
dz/dx這個中間值當然要保留下來以用於計算，所以粗略估計，backward的時候中間變量的佔用了是forward的兩倍！
優化器和動量
要注意，優化器也會佔用我們的顯存！
為什麼,看這個式子:
上式是典型的SGD隨機下降法的總體公式，權重W在進行更新的時候，會產生保存中間變量
，也就是在優化的時候，模型中的params參數所佔用的顯存量會翻倍。
當然這只是SGD優化器，其他複雜的優化器如果在計算時需要的中間變量多的時候，就會佔用更多的內存。
模型中哪些層會佔用顯存
有參數的層即會佔用顯存的層。我們一般的卷積層都會佔用顯存，而我們經常使用的激活層Relu沒有參數就不會佔用了。
佔用顯存的層一般是：
- 卷積層，通常的conv2d
- 全連接層，也就是Linear層
- BatchNorm層
- Embedding層
而不佔用顯存的則是：
- 剛才說到的激活層Relu等
- 池化層
- Dropout層
具體計算方式：
- Conv2d(Cin, Cout, K): 參數數目：Cin × Cout × K × K
- Linear(M->N): 參數數目：M×N
- BatchNorm(N): 參數數目： 2N
- Embedding(N,W): 參數數目： N × W
額外的顯存
總結一下，我們在總體的訓練中，佔用顯存大概分以下幾類：
- 模型中的參數(卷積層或其他有參數的層)
- 模型在計算時產生的中間參數(也就是輸入圖像在計算時每一層產生的輸入和輸出)
- backward的時候產生的額外的中間參數
- 優化器在優化時產生的額外的模型參數
但其實，我們佔用的顯存空間為什麼比我們理論計算的還要大，原因大概是因為深度學習框架一些額外的開銷吧，不過如果通過上面公式，理論計算出來的顯存和實際不會差太多的。
如何優化
優化除了算法層的優化，最基本的優化無非也就一下幾點：
- 減少輸入圖像的尺寸
- 減少batch，減少每次的輸入圖像數量
- 多使用下採樣，池化層
- 一些神經網絡層可以進行小優化，利用relu層中設置
inplace - 購買顯存更大的顯卡
- 從深度學習框架上面進行優化
撩我吧
- 如果你與我志同道合於此，老潘很願意與你交流；
- 如果你喜歡老潘的內容，歡迎關注和支持。
- 如果你喜歡我的文章，希望點贊👍 收藏 📁 評論 💬 三連一下~
想知道老潘是如何學習踩坑的，想與我交流問題~請關注公眾號「oldpan博客」。
老潘也會整理一些自己的私藏，希望能幫助到大家，點擊神秘傳送門獲取。
