前言
親，顯存炸了，你的顯卡快冒煙了！
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，這些是神經網路的權重大小，可以看到第一層卷積是3x3，而輸入影像的通道是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博客」，
老潘也會整理一些自己的私藏，希望能幫助到大家，點擊神秘傳送門獲取，
轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/264720.html
標籤：其他