前言
親，視訊記憶體炸了，你的顯示卡快冒煙了！
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部落格」。
老潘也會整理一些自己的私藏，希望能幫助到大家，點選神祕傳送門獲取。