淺談深度學習:如何計算模型以及中間變量的顯存佔用大小

前言

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

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博客」。
老潘也會整理一些本身的私藏，但願能幫助到你們，點擊神祕傳送門獲取。