PyTorch显存机制分析
參考：
=======================================================
在pytorch中有幾個關于顯存的關鍵詞：
在pytorch中顯存為緩存和變量分配的空間之和叫做reserved_memory，為變量分配的顯存叫做memory_allocated，由此可知reserved_memory一定大于等于memory_allocated，但是pytorch獲得總顯存要比reserved_memory要大，pytorch獲得的總顯存空間為reserved_memory+PyTorch context。
在不同顯卡和驅動下PyTorch context的大小是不同的，如：
https://zhuanlan.zhihu.com/p/424512257
所述，RTX 3090的context 開銷。其中3090用的CUDA 11.3，開銷為1639MB。
執行代碼：
import torch temp = torch.tensor([1.0]).cuda()
NVIDIA顯存消耗：
其中：
我們知道memory_reserved大小為2MB，那么context大小大約為1639MB。
給出
https://zhuanlan.zhihu.com/p/424512257
圖片：
可以知道，pytorch并沒有直接采用操作系統的顯存管理機制而是自己又寫了一個顯存管理機制，使用這種層級的管理機制在cache中申請顯存不需要向OS申請而是在自己的顯存管理程序中進行調配，如果自己的cache中顯存空間不夠再會通過OS來申請顯存，通過這種方法可以進一步提升顯存的申請速度和減少顯存碎片，當然這樣也有不好的地方，那就是多人使用共享顯卡的話容易導致一方一直不釋放顯存而另一方無法獲得足夠顯存，當然pytorch也給出了一些顯存分配的配置方法，但是主要還是為了減少顯存碎片的。
對https://zhuanlan.zhihu.com/p/424512257 中代碼進行一定修改：
import torch s = 0 # 模型初始化 linear1 = torch.nn.Linear(1024,1024, bias=False).cuda() # + 4194304 s = s+4194304 print(torch.cuda.memory_allocated(), s) linear2 = torch.nn.Linear(1024, 1, bias=False).cuda() # + 4096 s+=4096 print(torch.cuda.memory_allocated(), s) # 輸入定義 inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304 s+=4194304 print(torch.cuda.memory_allocated(), s) # 前向傳播 s=s+4194304+512 loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512 print(torch.cuda.memory_allocated(), s) # 后向傳播 loss.backward() # memory - 4194304 + 4194304 + 4096 s = s-4194304+4194304+4096 print(torch.cuda.memory_allocated(), s) # 再來一次~ loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 (512沒了，因為loss的ref還在) s+=4194304 print(torch.cuda.memory_allocated(), s) loss.backward() # memory - 4194304 s-=4194304 print(torch.cuda.memory_allocated(), s)
============================================
=================================================
修改代碼:
import torch
s = 0
# 模型初始化
linear1 = torch.nn.Linear(1024,1024, bias=False).cuda() # + 4194304
s = s+4194304
print(torch.cuda.memory_allocated(), s)
linear2 = torch.nn.Linear(1024, 1, bias=False).cuda() # + 4096
s+=4096
print(torch.cuda.memory_allocated(), s)
# 輸入定義
inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304
s+=4194304
print(torch.cuda.memory_allocated(), s)
# 前向傳播
s=s+4194304+512
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512
print(torch.cuda.memory_allocated(), s)
# 后向傳播
loss.backward() # memory - 4194304 + 4194304 + 4096
s = s-4194304+4194304+4096
print(torch.cuda.memory_allocated(), s)
# 再來一次~
for _ in range(10000):
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 (512沒了，因為loss的ref還在)
loss.backward() # memory - 4194304
print(torch.cuda.max_memory_reserved()/1024/1024, "MB")
print(torch.cuda.max_memory_allocated()/1024/1024, "MB")
print(torch.cuda.max_memory_cached()/1024/1024, "MB")
print(torch.cuda.memory_summary())
View Code
那么問題來了，問了保證這個程序完整運行下來的顯存量是多少呢？？？
已經知道最大的reserved_memory 為 22MB，那么保證該程序運行的最大顯存空間為reserved_memory+context_memory，
這里我們是使用1060G顯卡運行，先對一下context_memory:
執行代碼：
import torch temp = torch.tensor([1.0]).cuda()
NVIDIA顯存消耗：
所以context_memory為681MB-2MB=679MB
由于max_reserved_memory=22MB，因此該程序完整運行下來最高需要679+22=701MB，驗證一下：
再次運行代碼：
import torch
import time
s = 0
# 模型初始化
linear1 = torch.nn.Linear(1024,1024, bias=False).cuda() # + 4194304
s = s+4194304
print(torch.cuda.memory_allocated(), s)
linear2 = torch.nn.Linear(1024, 1, bias=False).cuda() # + 4096
s+=4096
print(torch.cuda.memory_allocated(), s)
# 輸入定義
inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304
s+=4194304
print(torch.cuda.memory_allocated(), s)
# 前向傳播
s=s+4194304+512
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512
print(torch.cuda.memory_allocated(), s)
# 后向傳播
loss.backward() # memory - 4194304 + 4194304 + 4096
s = s-4194304+4194304+4096
print(torch.cuda.memory_allocated(), s)
# 再來一次~
for _ in range(10000):
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 (512沒了，因為loss的ref還在)
loss.backward() # memory - 4194304
print(torch.cuda.max_memory_reserved()/1024/1024, "MB")
print(torch.cuda.max_memory_allocated()/1024/1024, "MB")
print(torch.cuda.max_memory_cached()/1024/1024, "MB")
print(torch.cuda.memory_summary())
time.sleep(60)
View Code
發現 803-701=102MB，這中間差的數值無法解釋，只能說memory_context可以隨著程序不同數值也不同，不同程序引入的pytorch函數不同導致context_memory也不同，這里我們按照這個想法反推，context_memory在這里為803-22=781MB，為了驗證我們修改代碼：
修改代碼：
import torch
import time
s = 0
# 模型初始化
linear1 = torch.nn.Linear(1024,1024*2, bias=False).cuda() # + 4194304
s = s+4194304
print(torch.cuda.memory_allocated(), s)
linear2 = torch.nn.Linear(1024*2, 1, bias=False).cuda() # + 4096
s+=4096
print(torch.cuda.memory_allocated(), s)
# 輸入定義
inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304
s+=4194304
print(torch.cuda.memory_allocated(), s)
# 前向傳播
s=s+4194304+512
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512
print(torch.cuda.memory_allocated(), s)
# 后向傳播
loss.backward() # memory - 4194304 + 4194304 + 4096
s = s-4194304+4194304+4096
print(torch.cuda.memory_allocated(), s)
# 再來一次~
for _ in range(100):
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 (512沒了，因為loss的ref還在)
loss.backward() # memory - 4194304
print(torch.cuda.max_memory_reserved()/1024/1024, "MB")
print(torch.cuda.max_memory_allocated()/1024/1024, "MB")
print(torch.cuda.max_memory_cached()/1024/1024, "MB")
print(torch.cuda.memory_summary())
time.sleep(60)
View Code
運行結果：
那么該代碼完整運行需要的顯存空間為：781+42=823MB
參考NVIDIA顯卡的顯存消耗：
發現支持剛才的猜想，也就是說不同的pytorch函數，顯卡型號，驅動，操作系統，cuda版本都是會影響context_memory大小的。
其中最為難以測定的就是pytorch函數，因為你可能一直在同一個平臺上跑代碼但是不太可能一直都用相同的pytorch函數，所以一個程序跑完最低需要的顯存空間的測定其實是需要完整跑一次網絡的反傳才可以測定的。
我這里采用的測定最低需要的顯存空間的方法是不考慮context_memory而去直接考慮一次反傳后最大需要的顯存，此時我們可以一次反傳后把程序掛住，如sleep一下，然后看下NVIDIA顯卡一共消耗了多少顯存。而且由上面的信息可知context_memory的測定是與具體使用的函數相關的，因此最穩妥的方法就是使用NVIDIA-smi監測一次完整反傳后最大顯存的消耗。
=====================================================
本博客是博主個人學習時的一些記錄，不保證是為原創，個別文章加入了轉載的源地址還有個別文章是匯總網上多份資料所成，在這之中也必有疏漏未加標注者，如有侵權請與博主聯系。
總結
以上是生活随笔為你收集整理的PyTorch显存机制分析的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 电脑经常自动重启的原因
- 下一篇： 参公