文 | Connolly@知乎（已授權）
源 |?極市平臺

作者最近兩年在研究分布式并行，經常使用PyTorch框架。一開始用的時候對于PyTorch的顯存機制也是一知半解，連蒙帶猜的，經常來知乎上來找答案，那么我就吸收大家的看法，為PyTorch的顯存機制做個小的總結吧。

實驗環境：
OS: Ubuntu18.04
python: 3.7.4
PyTorch: 1.9.1
GPU: V100

一、理論知識

1. 深度學習訓練過程

開門見山的說，PyTorch在進行深度學習訓練的時候，有4大部分的顯存開銷，分別是模型參數(parameters)，模型參數的梯度(gradients)，優化器狀態(optimizer states)以及中間激活值(intermediate activations) 或者叫中間結果(intermediate results)。

為了后面顯存分析闡述的方便，我將深度學習的訓練定義4個步驟：

1.模型定義：定義了模型的網絡結構，產生模型參數； while(你想訓練):2.前向傳播：執行模型的前向傳播，產生中間激活值；3.后向傳播：執行模型的后向傳播，產生梯度；4.梯度更新：執行模型參數的更新，第一次執行的時候產生優化器狀態。

在模型定義完之后，2~4循環執行。

2. 前向傳播

拿Linear層（或者叫Dense層，前饋神經網絡，全連接層等等...）舉例：假設他的權重矩陣為W，偏置向量為b，那么他的前向計算過程就是：

這里的X為該層的輸入向量，Y為輸出向量（中間激活值）

3. 后向傳播（反向傳播）

參考了這篇文章：

題目：《神經網絡反向傳播的數學原理》
鏈接：
https://zhuanlan.zhihu.com/p/22473137

后向傳播回來了一個第l+1層的輸出誤差矩陣 ，用以計算該層的梯度和輸入誤差

402 Payment Required

4. 梯度更新

接下來就是利用W_diff和b_diff進行更新了：

402 Payment Required

當然使用Adam優化器的時候，實際的更新過程并沒有上面的這么簡單。目前用的最多的是AdamW，可以看看這篇文章。

但是使用這一類優化器，也會帶來額外的顯存開銷。對于每一個參數，Adam都會為它準備對應的2個優化器狀態，分別是動量(momentum)和方差(variance)，用以加速模型的訓練。

二、顯存分析與Torch機制

2.1. 分析方法

(1) No Nvidia-smi

我看很多人現在還在用nvidia-smi來看pytorch的顯存占用，盯著跳來跳去的torch緩存區分析真的不累嗎。

而且PyTorch是有緩存區的設置的，意思就是一個Tensor就算被釋放了，進程也不會把空閑出來的顯存還給GPU，而是等待下一個Tensor來填入這一片被釋放的空間。

有什么好處？進程不需要重新向GPU申請顯存了，運行速度會快很多，有什么壞處？

他不能準確地給出某一個時間點具體的Tensor占用的顯存，而是顯示的已經分配到的顯存緩沖區和torch在創建cuda進程時所需開銷的和，也就是reserved_memory和torch context顯存之和。

這也是令很多人在使用PyTorch時對顯存占用感到困惑的罪魁禍首。

(2) torch.cuda is all you need.

在分析PyTorch的顯存時候，一定要使用torch.cuda里的顯存分析函數，我用的最多的是torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()，前者可以精準地反饋當前進程中torch.Tensor所占用的GPU顯存(注意是只包括torch.Tensor)，后者則可以告訴我們到調用函數為止所達到的最大的顯存占用字節數。

還有像torch.cuda.memory_reserved()這樣的函數則是查看當前進程所分配的顯存緩沖區是多少的。

非~常~好~用

▲Torch 官方文檔

2.2. PyTorch context開銷

-----之前沒有提到PyTorch context的開銷，做個補充...

我注意到有很多同學在做顯存分析的時候是為了在訓練的時候可以把卡的顯存用滿，這個之前沒有考慮到呢。其實PyTorch context是我們在使用torch的時候的一個大頭開銷。

主要參考的是論壇里的這篇討論：

How do I create Torch Tensor without any wasted storage space/baggage?

https://link.zhihu.com/?target=https%3A//discuss.pytorch.org/t/how-do-i-create-torch-tensor-without-any-wasted-storage-space-baggage/131134

什么是PyTorch context? 其實官方給他的稱呼是CUDA context，就是在第一次執行CUDA操作，也就是使用GPU的時候所需要創建的維護設備間工作的一些相關信息。如下圖所示

這個值跟CUDA的版本，pytorch的版本以及所使用的設備都是有關系的。目前我在ubuntu的torch1.9上測過RTX 3090和V100的context 開銷。其中3090用的CUDA 11.4，開銷為1639MB；V100用的CUDA 10.2，開銷為1351MB。

感興趣的同學可以在shell中執行下面這兩行代碼，然后用nvidia-smi去看看自己的環境里context的大小。然后用總大小減去context的大小再做顯存分析。

import torch temp = torch.tensor([1.0]).cuda()

我估計會有人問怎么去減小這個開銷...官方也給了一個辦法，看看自己有哪些cuda依賴是不需要的，比如cuDNN，然后自己重新編譯一遍PyTorch。編譯的時候把對應的包的flag給設為false就好了。我是還沒有試過，要搭編譯的環境太難受了，而且還要經常和庫做更新。

2.3 Torch顯存分配機制

在PyTorch中，顯存是按頁為單位進行分配的，這可能是CUDA設備的限制。就算我們只想申請4字節的顯存，CUDA也會為我們分配512字節或者1024字節的空間。

2.4?Torch顯存釋放機制

在PyTorch中，只要一個Tensor對象在后續不會再被使用，那么PyTorch就會自動回收該Tensor所占用的顯存，并以緩沖區的形式繼續占用顯存。要是實在看緩沖區不爽的話，也可以用torch.cuda.empty_cache()把它歸零，但是程序速度會變慢哦

3 訓練過程顯存分析

為了讓大家方便理解，我這里用torch.nn.Linear(1024, 1024, bias=False) 來做例子。

為了省事，loss函數則直接對輸出的樣本進行求和得到。沒辦法，想直接執行loss.backward()的話，loss得是標量才行呢。示例代碼：

import?torchmodel?=?torch.nn.Linear(1024,1024,?bias=False).cuda()? optimizer?=?torch.optim.AdamW(model.parameters()) inputs?=?torch.tensor([1.0]*1024).cuda()?#?shape?=?(1024) outputs?=?model(inputs)?#?shape?=?(1024) loss?=?sum(outputs)?#?shape?=?(1) loss.backward() optimizer.step()

3.1 模型的定義

結論：顯存占用量約為參數量乘以4

import?torchmodel?=?torch.nn.Linear(1024,1024,?bias=False).cuda()? print(torch.cuda.memory_allocated())

打印出來的數值為4194304，剛好等于1024×1024×4。

3.2 前向傳播過程

結論：顯存增加等于每一層模型產生的結果的顯存之和，且跟batch_size成正比。

inputs?=?torch.tensor([1.0]*1024).cuda()?#?shape?=?(1024)??memory?+?4096 outputs?=?model(inputs)?#?memory?+?4096

代碼中，outputs為產生的中間激活值，同時它也恰好是該模型的輸出結果。在執行完這一步之后，顯存增加了4096字節。(不算inputs的顯存的話)。

3.3 后向傳播過程

后向傳播會將模型的中間激活值給消耗并釋放掉掉，并為每一個模型中的參數計算其對應的梯度。在第一次執行的時候，會為模型參數分配對應的用來存儲梯度的空間。

loss?=?sum(outputs)?#?memory?+?512(torch?cuda分配最小單位) temp?=?torch.cuda.memory_allocated() loss.backward() print(torch.cuda.memory_allocated()?-?temp)?#?第一次增加4194304

第一次執行時顯存增加：4194304字節 - 激活值大小；第二次以后執行顯存減少：激活值大小；Note：由于這個中間激活值被賦給了outputs，所以后面在后向傳播的時候會發現，這個outputs的顯存沒有被釋放掉。

但是當層數變深的時候，就能明顯看到變化了。為了讓大家看到變化，再寫一段代碼~

import?torch#?模型初始化 linear1?=?torch.nn.Linear(1024,1024,?bias=False).cuda()?#?+?4194304 print(torch.cuda.memory_allocated()) linear2?=?torch.nn.Linear(1024,?1,?bias=False).cuda()?#?+?4096 print(torch.cuda.memory_allocated())#?輸入定義 inputs?=?torch.tensor([[1.0]*1024]*1024).cuda()?#?shape?=?(1024,1024)?#?+?4194304 print(torch.cuda.memory_allocated())#?前向傳播 loss?=?sum(linear2(linear1(inputs)))?#?shape?=?(1)?#?memory?+?4194304?+?512 print(torch.cuda.memory_allocated())#?后向傳播 loss.backward()?#?memory?-?4194304?+?4194304?+?4096 print(torch.cuda.memory_allocated())#?再來一次~ loss?=?sum(linear2(linear1(inputs)))?#?shape?=?(1)?#?memory?+?4194304??(512沒了，因為loss的ref還在) print(torch.cuda.memory_allocated()) loss.backward()?#?memory?-?4194304 print(torch.cuda.memory_allocated())

3.4 參數更新

optimizer.step()?#?第一次增加8388608，第二次就不增不減了哦

第一次執行時，會為每一個參數初始化其優化器狀態，對于這里的AdamW而言，每一個參數需要4*2=8個字節。

第二次開始，不會再額外分配顯存。顯存開銷:第一次: 增加8388608字節第二次及以后: 無增減3.5 Note由于計算機計算的特性，有一些計算操作在計算過程中是會帶來額外的顯存開銷的。

但是這種開銷在torch.memory_allocated中是不能被察覺的。

比如在AdamW在進行某一層的更新的時候，會帶來2倍該層參數量大小的臨時額外開銷。這個在max_memory_allocated中可以看到。在本例中就是8388608字節。

后臺回復關鍵詞【入群】

加入賣萌屋NLP、CV、搜廣推與求職討論群

[1]http://jwc.hnu.edu.cn/info/1185/10103.htm

[2]https://www.zhihu.com/question/547254904

[3]https://www.thepaper.cn/newsDetail_forward_19313268

總結

以上是生活随笔為你收集整理的Pytorch显存分配机制与显存占用分析方法的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。