本篇文章我們將介紹兩個 Python 庫 memory_profiler和Pytorch-Memory-Utils這兩個庫可以幫助我們了解內(nèi)存和顯存的消耗。

[[441124]]

memory_profiler

pip install memory_profiler#Load its magic function 
%load_ext memory_profiler 
from memory_profiler import profile 

memory_profiler可以完成以下的工作：

1、查找一行的內(nèi)存消耗

我們只需要在代碼的前面加上魔法函數(shù) %memit

%memit x = 10+5 
#Output 
peak memory: 54.01 MiB, increment: 0.27 MiB 

這里，峰值內(nèi)存(peak memory)是運行此代碼的進(jìn)程消耗的內(nèi)存。增量只是由于添加這行代碼而需要/消耗的內(nèi)存。同樣的邏輯也適用于以下其他的顯示。

2、查找函數(shù)的內(nèi)存消耗

在調(diào)用函數(shù)的行的開頭添加魔法函數(shù)。

def addition(): 
    a = [1] * (10 ** 1) 
    b = [2] * (3 * 10 ** 2) 
    sum = a+b 
    return sum 
 
%memit addition() 
#Output 
peak memory: 36.36 MiB, increment: 0.01 MiB 

3、逐行查找函數(shù)的內(nèi)存消耗

如果需要記錄函數(shù)中每一行的內(nèi)存使用，我們可以使用@profile 裝飾器。 但是@profile 僅適用于在單獨模塊中定義的函數(shù)，因此我們將首先使用 %%file 創(chuàng)建一個名為 demo.py 的簡單模塊，其中包含我們的函數(shù)

%%file demo.py 
from memory_profiler import profile 
 
@profile 
def addition(): 
    a = [1] * (10 ** 1) 
    b = [2] * (3 * 10 ** 2) 
    sum = a+b 
    return sum 

現(xiàn)在，我們可以調(diào)用該函數(shù)

from demo import addition 
 
%memit addition() 
 
#Output 
Line #    Mem usage    Increment   Line Contents 
================================================ 
     2     36.4 MiB     36.4 MiB   @profile 
     3                             def addition(): 
     4     36.4 MiB      0.0 MiB       a = [1] * (10 ** 1) 
     5   3851.1 MiB   3814.7 MiB       b = [2] * (3 * 10 ** 2) 
     6   7665.9 MiB   3814.8 MiB       sum = a+b 
     7   7665.9 MiB      0.0 MiB       return sum 
peak memory: 7665.88 MiB, increment: 7629.52 MiB 

4、完整python腳本的內(nèi)存消耗

這個這個方法不能再 notebook 中使用。我們必須創(chuàng)建 python 腳本并通過命令行運行它。

#create script.py 
import time 
 
@profile 
def function1(): 
    n = 100000 
    a = [1] * n 
    time.sleep(1) 
    return a 
     
@profile 
def function2(): 
    n = 200000 
    b = [1] * n 
    time.sleep(1) 
    return b 
 
if __name__ == "__main__": 
    function1() 
    function2() 

之后運行腳本并查看

#On command line 
mprof run script.py 
#To generate plot  
mprof plot 

我們可以看到內(nèi)存消耗與時間的關(guān)系圖

@profile裝飾器沒有必要放在函數(shù)前面，如果我們不保留它，我們不會看到函數(shù)級內(nèi)存消耗，但我們會看到整個腳本的內(nèi)存消耗

Pytorch-Memory-Utils

通過Pytorch-Memory-Utils工具，我們在使用顯存的代碼中間插入檢測函數(shù)，這樣就可以輸出在當(dāng)前行代碼時所占用的顯存。這個對于我們計算模型的GPU顯存占用是非常方便的，通過計算顯存占用，我們才能夠最大化訓(xùn)練的batch size，保證訓(xùn)練的最優(yōu)速度。

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)      # 創(chuàng)建顯存檢測對象 
 
gpu_tracker.track()                  # 開始檢測 
cnn = models.vgg19(pretrained=True).to(device)  # 導(dǎo)入VGG19模型并且將數(shù)據(jù)轉(zhuǎn)到顯存中 
gpu_tracker.track() 

然后可以發(fā)現(xiàn)程序運行過程中的顯存變化(第一行是載入前的顯存，最后一行是載入后的顯存)：

At __main__ <module>: line 13                        Total Used Memory:472.2  Mb 
 
+ | 1 * Size:(128, 64, 3, 3)      | Memory: 0.2949 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(256, 128, 3, 3)     | Memory: 1.1796 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(64, 64, 3, 3)       | Memory: 0.1474 M | <class 'torch.nn.parameter.Parameter'> 
+ | 2 * Size:(4096,)              | Memory: 0.0327 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(512, 256, 3, 3)     | Memory: 4.7185 M | <class 'torch.nn.parameter.Parameter'> 
+ | 2 * Size:(128,)               | Memory: 0.0010 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(1000, 4096)         | Memory: 16.384 M | <class 'torch.nn.parameter.Parameter'> 
+ | 6 * Size:(512,)               | Memory: 0.0122 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(64, 3, 3, 3)        | Memory: 0.0069 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(4096, 25088)        | Memory: 411.04 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(4096, 4096)         | Memory: 67.108 M | <class 'torch.nn.parameter.Parameter'> 
+ | 5 * Size:(512, 512, 3, 3)     | Memory: 47.185 M | <class 'torch.nn.parameter.Parameter'> 
+ | 2 * Size:(64,)                | Memory: 0.0005 M | <class 'torch.nn.parameter.Parameter'> 
+ | 3 * Size:(256,)               | Memory: 0.0030 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(128, 128, 3, 3)     | Memory: 0.5898 M | <class 'torch.nn.parameter.Parameter'> 
+ | 2 * Size:(256, 256, 3, 3)     | Memory: 4.7185 M | <class 'torch.nn.parameter.Parameter'> 
+ | 1 * Size:(1000,)              | Memory: 0.004 M | <class 'torch.nn.parameter.Parameter'> 
 
At __main__ <module>: line 15                        Total Used Memory:1387.5 Mb 

通過上面的報告，很容易發(fā)現(xiàn)一個問題。

首先我們知道VGG19所有層的權(quán)重大小加起來大約是548M(這個數(shù)值來源于Pytorch官方提供的VGG19權(quán)重文件大小)，我們將上面報告打印的Tensor-Memory也都加起來算下來也差不多551.8Mb。但是，我們算了兩次打印的顯存實際占用中：1387.5 – 472.2 = 915.3 MB。

Pytorch在開始運行程序時需要額外的顯存開銷，這種額外的顯存開銷與我們實際使用的模型權(quán)重顯存大小無關(guān)。 這個額外的顯存Pytorch的開發(fā)者也對此進(jìn)行說明了，這部分釋放后的顯存可以用，只不過不在Nvidia-smi中顯示罷了，所以我們無需關(guān)注。