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

01 理論知識

1.1 深度學習訓練過程

開門見山的說，PyTorch在進行深度學習訓練的時候，有4大部分的顯存開銷，分別是模型參數(shù)(parameters)，模型參數(shù)的梯度(gradients)，優(yōu)化器狀態(tài)(optimizer states)以及中間激活值(intermediate activations) 或者叫中間結(jié)果(intermediate results)。為了后面顯存分析闡述的方便，我將深度學習的訓練定義4個步驟：

1. 模型定義：定義了模型的網(wǎng)絡結(jié)構，產(chǎn)生模型參數(shù)；

while(你想訓練):

2. 前向傳播：執(zhí)行模型的前向傳播，產(chǎn)生中間激活值；
3. 后向傳播：執(zhí)行模型的后向傳播，產(chǎn)生梯度；
4. 梯度更新：執(zhí)行模型參數(shù)的更新，第一次執(zhí)行的時候產(chǎn)生優(yōu)化器狀態(tài)。

在模型定義完之后，2~4循環(huán)執(zhí)行。

1.2 前向傳播

拿Linear層（或者叫Dense層，前饋神經(jīng)網(wǎng)絡，全連接層等等...）舉例：假設他的權重矩陣為W，偏置向量為b，那么他的前向計算過程就是：，這里的X為該層的輸入向量，Y為輸出向量（中間激活值）

1.3 后向傳播（反向傳播）

參考了這篇文章《神經(jīng)網(wǎng)絡反向傳播的數(shù)學原理》https://zhuanlan.zhihu.com/p/22473137后向傳播回來了一個第l+1層的輸出誤差矩陣，用以計算該層的梯度和輸入誤差

1.4 梯度更新

接下來就是利用 W_diff 和 b_diff 進行更新了： 當然使用 Adam 優(yōu)化器的時候，實際的更新過程并沒有上面的這么簡單。目前用的最多的是 AdamW ，可以看看這篇文章《當前訓練神經(jīng)網(wǎng)絡最快的方式：AdamW優(yōu)化算法+超級收斂》https://zhuanlan.zhihu.com/p/38945390)但是使用這一類優(yōu)化器，也會帶來額外的顯存開銷。對于每一個參數(shù)，Adam都會為它準備對應的2個優(yōu)化器狀態(tài)，分別是動量(momentum)和方差(variance)，用以加速模型的訓練。02 顯存分析方法與Torch機制

2.1 分析方法

(1) No Nvidia-smi我看很多人現(xiàn)在還在用 nvidia-smi 來看 pytorch 的顯存占用，盯著跳來跳去的torch緩存區(qū)分析真的不累嗎。（貼一個Torch為什么不用Nvidia-smi看的圖）。而且PyTorch是有緩存區(qū)的設置的，意思就是一個Tensor就算被釋放了，進程也不會把空閑出來的顯存還給GPU，而是等待下一個Tensor來填入這一片被釋放的空間。有什么好處？進程不需要重新向GPU申請顯存了，運行速度會快很多，有什么壞處？他不能準確地給出某一個時間點具體的Tensor占用的顯存，而是顯示的已經(jīng)分配到的顯存和顯存緩沖區(qū)之和。這也是令很多人在使用PyTorch時對顯存占用感到困惑的罪魁禍首。(2) torch.cuda is all you need在分析PyTorch的顯存時候，一定要使用torch.cuda里的顯存分析函數(shù)，我用的最多的是torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()，前者可以精準地反饋當前進程中Torch.Tensor所占用的GPU顯存，后者則可以告訴我們到調(diào)用函數(shù)為止所達到的最大的顯存占用字節(jié)數(shù)。還有像torch.cuda.memory_reserved()這樣的函數(shù)則是查看當前進程所分配的顯存緩沖區(qū)是多少的。memory_allocated+memory_reserved就等于nvidia-smi中的值啦。非~常~好~用chao dashengTorch 官方文檔2.2 PyTorch context開銷-----之前沒有提到PyTorch context的開銷，做個補充...我注意到有很多同學在做顯存分析的時候是為了在訓練的時候可以把卡的顯存用滿，這個之前沒有考慮到呢。其實PyTorch context是我們在使用torch的時候的一個大頭開銷。主要參考的是論壇里的這篇討論：How do I create Torch Tensor without any wasted storage space/baggage?https://discuss.pytorch.org/t/how-do-i-create-torch-tensor-without-any-wasted-storage-space-baggage/131134什么是PyTorch context? 其實官方給他的稱呼是CUDA context，就是在第一次執(zhí)行CUDA操作，也就是使用GPU的時候所需要創(chuàng)建的維護設備間工作的一些相關信息。如下圖所示這個值跟CUDA的版本，pytorch的版本以及所使用的設備都是有關系的。目前我在ubuntu的torch1.9上測過RTX 3090和V100的context 開銷。其中3090用的CUDA 11.4，開銷為1639MB；V100用的CUDA 10.2，開銷為1351MB。感興趣的同學可以在shell中執(zhí)行下面這兩行代碼，然后用nvidia-smi去看看自己的環(huán)境里context的大小。然后用總大小減去context的大小再做顯存分析。

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

2.3Torch顯存分配機制

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

2.4Torch顯存釋放機制

在PyTorch中，只要一個Tensor對象在后續(xù)不會再被使用，那么PyTorch就會自動回收該Tensor所占用的顯存，并以緩沖區(qū)的形式繼續(xù)占用顯存。要是實在看緩沖區(qū)不爽的話，也可以用torch.cuda.empty_cache()把它歸零，但是程序速度會變慢哦03 訓練過程顯存分析為了讓大家方便理解，我這里用torch.nn.Linear(1024, 1024, bias=False) 來做例子。為了省事，loss函數(shù)則直接對輸出的樣本進行求和得到。沒辦法，想直接執(zhí)行l(wèi)oss.backward()的話，loss得是標量才行呢。示例代碼：

import torch model = 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 模型的定義

結(jié)論：顯存占用量約為參數(shù)量乘以4

import torch model = torch.nn.Linear(1024,1024, bias=False).cuda() print(torch.cuda.memory_allocated())打印出來的數(shù)值為4194304，剛好等于1024×1024×4。

3.2 前向傳播過程

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

inputs = torch.tensor([1.0]*1024).cuda() # shape = (1024) memory + 4096 outputs = model(inputs) # memory + 4096代碼中，outputs為產(chǎn)生的中間激活值，同時它也恰好是該模型的輸出結(jié)果。在執(zhí)行完這一步之后，顯存增加了4096字節(jié)。(不算inputs的顯存的話)。

3.3 后向傳播過程

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

loss = sum(outputs) # memory + 512(torch cuda分配最小單位) temp = torch.cuda.memory_allocated() loss.backward() print(torch.cuda.memory_allocated() - temp) # 第一次增加4194304第一次執(zhí)行時顯存增加：4194304字節(jié) - 激活值大??；第二次以后執(zhí)行顯存減少：激活值大小；Note：由于這個中間激活值被賦給了outputs，所以后面在后向傳播的時候會發(fā)現(xiàn)，這個outputs的顯存沒有被釋放掉。但是當層數(shù)變深的時候，就能明顯看到變化了。為了讓大家看到變化，再寫一段代碼~

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 參數(shù)更新

optimizer.step()#第一次增加8388608，第二次就不增不減了哦第一次執(zhí)行時，會為每一個參數(shù)初始化其優(yōu)化器狀態(tài)，對于這里的AdamW而言，每一個參數(shù)需要4*2=8個字節(jié)。第二次開始，不會再額外分配顯存。顯存開銷:第一次: 增加8388608字節(jié)第二次及以后: 無增減3.5 Note由于計算機計算的特性，有一些計算操作在計算過程中是會帶來額外的顯存開銷的。但是這種開銷在torch.memory_allocated中是不能被察覺的。比如在AdamW在進行某一層的更新的時候，會帶來2倍該層參數(shù)量大小的臨時額外開銷。這個在max_memory_allocated中可以看到。在本例中就是8388608字節(jié)。

審核編輯 ：李倩