# 前言 親，視訊記憶體炸了，你的顯示卡快冒煙了！ ```python 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` - 購買視訊記憶體更大的顯示卡 - 從深度學習框架上面進行優化 # 撩我吧 - 如果你與我志同道合於此，老潘很願意與你交流； - 如果你喜歡老潘的內容，歡迎關注和支援。 - 如果你喜歡我的文章，希望點贊