释放pytorch占用的gpu显存_再次浅谈Pytorch中的显存利用问题(附完善显存跟踪代码)...
前言
之前在淺談深度學習:如何計算模型以及中間變量的顯存占用大小和如何在Pytorch中精細化利用顯存中我們已經談論過了平時使用中顯存的占用來自于哪里，以及如何在Pytorch中更好地使用顯存。在這篇文章中，我們借用Pytorch-Memory-Utils這個工具來檢測我們在訓練過程中關于顯存的變化情況，分析出我們如何正確釋放多余的顯存。
在深度探究前先了解下我們的輸出信息，通過Pytorch-Memory-Utils工具，我們在使用顯存的代碼中間插入檢測函數(如何使用見工具github頁面和下文部分)，就可以輸出類似于下面的信息，At __main__ : line 13 Total Used Memory:696.5 Mb表示在當前行代碼時所占用的顯存，即在我們的代碼中執行到13行的時候所占顯存為695.5Mb。At __main__ : line 15 Total Used Memory:1142.0 Mb表示程序執行到15行時所占的顯存為1142.0Mb。兩條數據之間表示所占顯存的tensor變量。
# 12-Sep-18-21:48:45-gpu_mem_track.txt
GPU Memory Track | 12-Sep-18-21:48:45 | Total Used Memory:696.5 Mb
At __main__ : line 13 Total Used Memory:696.5 Mb
+ | 7 * Size:(512, 512, 3, 3) | Memory: 66.060 M |
+ | 1 * Size:(512, 256, 3, 3) | Memory: 4.7185 M |
+ | 1 * Size:(64, 64, 3, 3) | Memory: 0.1474 M |
+ | 1 * Size:(128, 64, 3, 3) | Memory: 0.2949 M |
+ | 1 * Size:(128, 128, 3, 3) | Memory: 0.5898 M |
+ | 8 * Size:(512,) | Memory: 0.0163 M |
+ | 3 * Size:(256, 256, 3, 3) | Memory: 7.0778 M |
+ | 1 * Size:(256, 128, 3, 3) | Memory: 1.1796 M |
+ | 2 * Size:(64,) | Memory: 0.0005 M |
+ | 4 * Size:(256,) | Memory: 0.0040 M |
+ | 2 * Size:(128,) | Memory: 0.0010 M |
+ | 1 * Size:(64, 3, 3, 3) | Memory: 0.0069 M |
At __main__ : line 15 Total Used Memory:1142.0 Mb
+ | 1 * Size:(60, 3, 512, 512) | Memory: 188.74 M |
+ | 1 * Size:(30, 3, 512, 512) | Memory: 94.371 M |
+ | 1 * Size:(40, 3, 512, 512) | Memory: 125.82 M |
At __main__ : line 21 Total Used Memory:1550.9 Mb
+ | 1 * Size:(120, 3, 512, 512) | Memory: 377.48 M |
+ | 1 * Size:(80, 3, 512, 512) | Memory: 251.65 M |
At __main__ : line 26 Total Used Memory:2180.1 Mb
- | 1 * Size:(120, 3, 512, 512) | Memory: 377.48 M |
- | 1 * Size:(40, 3, 512, 512) | Memory: 125.82 M |
At __main__ : line 32 Total Used Memory:1676.8 Mb
當然這個檢測工具不僅適用于Pytorch，其他的深度學習框架也同樣可以使用，不過需要注意下靜態圖和動態圖在實際運行過程中的區別。
正文
了解了Pytorch-Memory-Utils工具如何使用后，接下來我們通過若干段程序代碼來演示在Pytorch訓練中：
平時的顯存是如何變化的，到底是什么占用了顯存。
如何去釋放不需要的顯存。
首先，我們在下段代碼中導入我們需要的庫，隨后開始我們的顯存檢測程序。
import torch
import inspect
from torchvision import models
from gpu_mem_track import MemTracker # 引用顯存跟蹤代碼
device = torch.device('cuda:0')
frame = inspect.currentframe()
gpu_tracker = MemTracker(frame) # 創建顯存檢測對象
gpu_tracker.track() # 開始檢測
預訓練權重模型
首先我們檢測一下神經網絡模型權重所占用的顯存信息，下面代碼中我們嘗試加載VGG19這個經典的網絡模型，并且導入預訓練好的權重。
gpu_tracker.track()
cnn = models.vgg19(pretrained=True).to(device) # 導入VGG19模型并且將數據轉到顯存中
gpu_tracker.track()
然后可以發現程序運行過程中的顯存變化(第一行是載入前的顯存，最后一行是載入后的顯存)：
At __main__ : line 13 Total Used Memory:472.2 Mb
+ | 1 * Size:(128, 64, 3, 3) | Memory: 0.2949 M |
+ | 1 * Size:(256, 128, 3, 3) | Memory: 1.1796 M |
+ | 1 * Size:(64, 64, 3, 3) | Memory: 0.1474 M |
+ | 2 * Size:(4096,) | Memory: 0.0327 M |
+ | 1 * Size:(512, 256, 3, 3) | Memory: 4.7185 M |
+ | 2 * Size:(128,) | Memory: 0.0010 M |
+ | 1 * Size:(1000, 4096) | Memory: 16.384 M |
+ | 6 * Size:(512,) | Memory: 0.0122 M |
+ | 1 * Size:(64, 3, 3, 3) | Memory: 0.0069 M |
+ | 1 * Size:(4096, 25088) | Memory: 411.04 M |
+ | 1 * Size:(4096, 4096) | Memory: 67.108 M |
+ | 5 * Size:(512, 512, 3, 3) | Memory: 47.185 M |
+ | 2 * Size:(64,) | Memory: 0.0005 M |
+ | 3 * Size:(256,) | Memory: 0.0030 M |
+ | 1 * Size:(128, 128, 3, 3) | Memory: 0.5898 M |
+ | 2 * Size:(256, 256, 3, 3) | Memory: 4.7185 M |
+ | 1 * Size:(1000,) | Memory: 0.004 M |
At __main__ : line 15 Total Used Memory:1387.5 Mb
通過上面的報告，很容易發現一個問題。
首先我們知道VGG19所有層的權重大小加起來大約是548M(這個數值來源于Pytorch官方提供的VGG19權重文件大小)，我們將上面報告打印的Tensor-Memory也都加起來算下來也差不多551.8Mb。但是，我們算了兩次打印的顯存實際占用中：1387.5 – 472.2 = 915.3 MB。
唉，怎么多用了差不多400Mb呢？是不是報告出什么問題了。
這樣，我們再加點Tensor試一下。
...
gpu_tracker.track()
cnn = models.vgg19(pretrained=True).to(device)
gpu_tracker.track()
# 上方為之前的代碼
# 新增加的tensor
dummy_tensor_1 = torch.randn(30, 3, 512, 512).float().to(device) # 30*3*512*512*4/1000/1000 = 94.37M
dummy_tensor_2 = torch.randn(40, 3, 512, 512).float().to(device) # 40*3*512*512*4/1000/1000 = 125.82M
dummy_tensor_3 = torch.randn(60, 3, 512, 512).float().to(device) # 60*3*512*512*4/1000/1000 = 188.74M
gpu_tracker.track() # 再次打印
如上面的代碼，我們又加入了三個Tensor，全部放到顯存中。報告如下：
At __main__ : line 15 Total Used Memory:1387.5 Mb
+ | 1 * Size:(30, 3, 512, 512) | Memory: 94.371 M |
+ | 1 * Size:(40, 3, 512, 512) | Memory: 125.82 M |
+ | 1 * Size:(60, 3, 512, 512) | Memory: 188.74 M |
At __main__ : line 21 Total Used Memory:1807.0 Mb
上面的報告就比較正常了：94.3 + 125.8 + 188.7 = 408.8 約等于 1807.0 – 1387.5 = 419.5，誤差可以忽略，因為肯定會存在一些開銷使用的顯存。
那之前是什么情況？是不是模型的權重信息占得顯存就稍微多一點？
這樣，我們將載入VGG19模型的代碼注釋掉，只對后面的三個Tensor進行檢測。
...
gpu_tracker.track()
# cnn = models.vgg19(pretrained=True).to(device) 注釋掉讀權重代碼
gpu_tracker.track()
...
可以發現：
GPU Memory Track | 15-Sep-18-13:59:03 | Total Used Memory:513.3 Mb
At __main__ : line 13 Total Used Memory:513.3 Mb
At __main__ : line 15 Total Used Memory:513.3 Mb
At __main__ : line 18 Total Used Memory:513.3 Mb
+ | 1 * Size:(60, 3, 512, 512) | Memory: 188.74 M |
+ | 1 * Size:(30, 3, 512, 512) | Memory: 94.371 M |
+ | 1 * Size:(40, 3, 512, 512) | Memory: 125.82 M |
At __main__ : line 24 Total Used Memory:1271.3 Mb
同樣，顯存占用比所列出來的Tensor占用大，我們暫時將次歸結為Pytorch在開始運行程序時需要額外的顯存開銷，這種額外的顯存開銷與我們實際使用的模型權重顯存大小無關。
Pytorch使用的顯存策略
Pytorch已經可以自動回收我們“不用的”顯存，類似于python的引用機制，當某一內存內的數據不再有任何變量引用時，這部分的內存便會被釋放。但有一點需要注意，當我們有一部分顯存不再使用的時候，這部分釋放的顯存通過Nvidia-smi命令是看不到的，舉個例子：
device = torch.device('cuda:0')
# 定義兩個tensor
dummy_tensor_4 = torch.randn(120, 3, 512, 512).float().to(device) # 120*3*512*512*4/1000/1000 = 377.48M
dummy_tensor_5 = torch.randn(80, 3, 512, 512).float().to(device) # 80*3*512*512*4/1000/1000 = 251.64M
# 然后釋放
dummy_tensor_4 = dummy_tensor_4.cpu()
dummy_tensor_2 = dummy_tensor_2.cpu()
# 這里雖然將上面的顯存釋放了，但是我們通過Nvidia-smi命令看到顯存依然在占用
torch.cuda.empty_cache()
# 只有執行完上面這句，顯存才會在Nvidia-smi中釋放
Pytorch的開發者也對此進行說明了，這部分釋放后的顯存可以用，只不過不在Nvidia-smi中顯示罷了。
關于模型調用
torch.no_grad()是Pytorch-0.4版本時候更新的功能，在此語句的作用域下，所有的tensor運算不會保存梯度值，特別適合在inference的時候使用，代替舊版本的volatile。
用一段代碼演示下，這里我們根據VGG19網絡構造一個特征提取器，分別提取content_image和style_image的特征圖，然后將提取的特征圖存在兩個list中,我們使用了with torch.no_grad()語句(在沒使用no_grad之前占用的顯存更多，不過這里不進行展示了)：
gpu_tracker.track()
layers = ['relu_1', 'relu_3', 'relu_5', 'relu_9'] # 提取的層數
layerIdx = 0
content_image = torch.randn(1, 3, 500, 500).float().to(device)
style_image = torch.randn(1, 3, 500, 500).float().to(device)
feature_extractor = nn.Sequential().to(device) # 特征提取器
cnn = models.vgg19(pretrained=True).features.to(device) # 采取VGG19
input_features = [] # 保存提取出的features
target_features = [] # 保存提取出的features
i = 0
# 如果不加下面這一句,那么顯存的占用提升,因為保存了中間計算的梯度值
with torch.no_grad():
for layer in cnn.children():
if layerIdx < len(layers):
if isinstance(layer, nn.Conv2d):
i += 1
name = "conv_" + str(i)
feature_extractor.add_module(name, layer)
elif isinstance(layer, nn.MaxPool2d):
name = "pool_" + str(i)
feature_extractor.add_module(name, layer)
elif isinstance(layer, nn.ReLU):
name = "relu_" + str(i)
feature_extractor.add_module(name, nn.ReLU(inplace=True))
if name == layers[layerIdx]:
input = feature_extractor(content_image)
gpu_tracker.track()
target = feature_extractor(style_image)
gpu_tracker.track()
input_features.append(input)
target_features.append(target)
del input
del target
layerIdx += 1
gpu_tracker.track()
進行GPU跟蹤后，觀察下顯存變化：
At __main__ : line 33 Total Used Memory:1313.3 Mb
+ | 2 * Size:(64,) | Memory: 0.0005 M |
+ | 2 * Size:(1, 3, 500, 500) | Memory: 6.0 M |
+ | 1 * Size:(64, 64, 3, 3) | Memory: 0.1474 M |
+ | 1 * Size:(128, 64, 3, 3) | Memory: 0.2949 M |
+ | 2 * Size:(128,) | Memory: 0.0010 M |
+ | 2 * Size:(1, 256, 125, 125) | Memory: 32.0 M |
+ | 1 * Size:(128, 128, 3, 3) | Memory: 0.5898 M |
+ | 7 * Size:(512, 512, 3, 3) | Memory: 66.060 M |
+ | 3 * Size:(256, 256, 3, 3) | Memory: 7.0778 M |
+ | 2 * Size:(1, 512, 62, 62) | Memory: 15.745 M |
+ | 1 * Size:(64, 3, 3, 3) | Memory: 0.0069 M |
+ | 2 * Size:(1, 128, 250, 250) | Memory: 64.0 M |
+ | 8 * Size:(512,) | Memory: 0.0163 M |
+ | 4 * Size:(256,) | Memory: 0.0040 M |
+ | 1 * Size:(256, 128, 3, 3) | Memory: 1.1796 M |
+ | 1 * Size:(512, 256, 3, 3) | Memory: 4.7185 M |
+ | 2 * Size:(1, 64, 500, 500) | Memory: 128.0 M |
At __main__ : line 76 Total Used Memory:1932.0 Mb
上表中4*2個是提取出的特征圖，其他的則是模型的權重值，但是發現，所有的值加起來，與總顯存變化又不同，那究竟多了哪些占用顯存的東西？
其實原因很簡單，除了在程序運行時的一些額外顯存開銷，另外一個占用顯存的東西就是我們在計算時候的臨時緩沖值，這些零零總總也會占用一部分顯存，并且這些緩沖值通過Python的垃圾收集是收集不到的。
Asynchronous execution
做過并行計算或者操作系統的同學可能知道，GPU的計算方式一般是異步的。異步運算不像同步運算那樣是按照順序一步一步來，異步是同時進行的，異步計算中，兩種不一樣的操作可能會發生同時觸發的情況，這是處理兩者間的前后關系、依賴關系或者沖突關系就比較重要了。
有一個眾所周知的小技巧，在執行訓練程序的時候將環境變量CUDA_LAUNCH_BLOCKING=1設為1(強制同步)可以準確定位觀察到我們顯存操作的錯誤代碼行數。
后記
暫時就說這些，Pytorch的顯存優化除了以上這些，更多的應該交給底層處理了，期待一下Pytorch的再次更新吧——另外，Pytorch-1.0的dev版已經出來，大家可以嘗嘗鮮了！
如果想進一步了解：
總結
以上是生活随笔為你收集整理的释放pytorch占用的gpu显存_再次浅谈Pytorch中的显存利用问题(附完善显存跟踪代码)...的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： requirejs+jquery表单验证
- 下一篇： 解决chm文件打开后跳到某一主页方法