前言

親，顯存炸了，你的顯卡快冒煙了！

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 ，顯然是顯存裝不下你那么多的模型權重還有中間變量，然后程序奔潰了。怎么辦，其實辦法有很多，及時清空中間變量，優(yōu)化代碼，減少batch，等等等等，都能夠減少顯存溢出的風險。

但是這篇要說的是上面這一切優(yōu)化操作的基礎，如何去計算我們所使用的顯存。學會如何計算出來我們設計的模型以及中間變量所占顯存的大小，想必知道了這一點，我們對自己顯存也就會得心應手了。

如何計算

首先我們應該了解一下基本的數(shù)據(jù)量信息：

• 1 G = 1000 MB
• 1 M = 1000 KB
• 1 K = 1000 Byte
• 1 B = 8 bit

好，肯定有人會問為什么是1000而不是1024，這里不過多討論，只能說兩種說法都是正確的，只是應用場景略有不同。這里統(tǒng)一按照上面的標準進行計算。

然后我們說一下我們平常使用的向量所占的空間大小，以Pytorch官方的數(shù)據(jù)格式為例（所有的深度學習框架數(shù)據(jù)格式都遵循同一個標準）：

我們只需要看左邊的信息，在平常的訓練中，我們經(jīng)常使用的一般是這兩種類型：

• float32 單精度浮點型
• int32 整型

一般一個8-bit的整型變量所占的空間為 1B 也就是 8bit 。而32位的float則占 4B 也就是 32bit 。而雙精度浮點型double和長整型long在平常的訓練中我們一般不會使用。

ps：消費級顯卡對單精度計算有優(yōu)化，服務器級別顯卡對雙精度計算有優(yōu)化。

也就是說，假設有一幅RGB三通道真彩色圖片，長寬分別為500 x 500，數(shù)據(jù)類型為單精度浮點型，那么這張圖所占的顯存的大小為：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)所占的空間并不大，那為什么我們的顯存依舊還是用的比較多，原因很簡單，占用顯存比較多空間的并不是我們輸入圖像，而是神經(jīng)網(wǎng)絡中的中間變量以及使用optimizer算法時產生的巨量的中間參數(shù)。

我們首先來簡單計算一下Vgg16這個net需要占用的顯存：

通常一個模型占用的顯存也就是兩部分：

• 模型自身的參數(shù)(params)
• 模型計算產生的中間變量(memory)

圖片來自cs231n，這是一個典型的sequential-net，自上而下很順暢，我們可以看到我們輸入的是一張224x224x3的三通道圖像，可以看到一張圖像只占用 150x4k ，但上面是 150k ，這是因為這里在計算的時候默認的數(shù)據(jù)格式是8-bit而不是32-bit，所以最后的結果要乘上一個4。

我們可以看到，左邊的memory值代表：圖像輸入進去，圖片以及所產生的中間卷積層所占的空間。我們都知道，這些形形色色的深層卷積層也就是深度神經(jīng)網(wǎng)絡進行“思考”的過程：

圖片從3通道變?yōu)?4 --> 128 --> 256 --> 512 .... 這些都是卷積層，而我們的顯存也主要是他們占用了。

還有上面右邊的params，這些是神經(jīng)網(wǎng)絡的權重大小，可以看到第一層卷積是3x3，而輸入圖像的通道是3，輸出通道是64，所以很顯然，第一個卷積層權重所占的空間是 (3 x 3 x 3) x 64。

另外還有一個需要注意的是中間變量在backward的時候會翻倍！

舉個例子，下面是一個計算圖，輸入 x ，經(jīng)過中間結果 z ，然后得到最終變量 L ：

我們在backward的時候需要保存下來的中間值。輸出是 L ，然后輸入 x ，我們在backward的時候要求 L 對 x 的梯度，這個時候就需要在計算鏈 L 和 x 中間的 z ：

dz/dx 這個中間值當然要保留下來以用于計算，所以粗略估計， backward 的時候中間變量的占用了是 forward 的兩倍！

優(yōu)化器和動量

要注意，優(yōu)化器也會占用我們的顯存！

為什么,看這個式子:

上式是典型的SGD隨機下降法的總體公式，權重 W 在進行更新的時候，會產生保存中間變量  ，也就是在優(yōu)化的時候，模型中的params參數(shù)所占用的顯存量會翻倍。

當然這只是SGD優(yōu)化器，其他復雜的優(yōu)化器如果在計算時需要的中間變量多的時候，就會占用更多的內存。

模型中哪些層會占用顯存

有參數(shù)的層即會占用顯存的層。我們一般的卷積層都會占用顯存，而我們經(jīng)常使用的激活層Relu沒有參數(shù)就不會占用了。

占用顯存的層一般是：

• 卷積層，通常的conv2d
• 全連接層，也就是Linear層
• BatchNorm層
• Embedding層

而不占用顯存的則是：

• 剛才說到的激活層Relu等
• 池化層
• Dropout層

具體計算方式：

• Conv2d(Cin, Cout, K): 參數(shù)數(shù)目：Cin × Cout × K × K
• Linear(M->N): 參數(shù)數(shù)目：M×N
• BatchNorm(N): 參數(shù)數(shù)目： 2N
• Embedding(N,W): 參數(shù)數(shù)目： N × W

額外的顯存

總結一下，我們在總體的訓練中，占用顯存大概分以下幾類：

• 模型中的參數(shù)(卷積層或其他有參數(shù)的層)
• 模型在計算時產生的中間參數(shù)(也就是輸入圖像在計算時每一層產生的輸入和輸出)
• backward的時候產生的額外的中間參數(shù)
• 優(yōu)化器在優(yōu)化時產生的額外的模型參數(shù)

但其實，我們占用的顯存空間為什么比我們理論計算的還要大，原因大概是因為深度學習框架一些額外的開銷吧，不過如果通過上面公式，理論計算出來的顯存和實際不會差太多的。

如何優(yōu)化

優(yōu)化除了算法層的優(yōu)化，最基本的優(yōu)化無非也就一下幾點：

inplace 

撩我吧

• 如果你與我志同道合于此，老潘很愿意與你交流；
• 如果你喜歡老潘的內容，歡迎關注和支持。
• 如果你喜歡我的文章，希望點贊:+1: 收藏 :file_folder: 評論 :speech_balloon: 三連一下~