乐趣区

关于人工智能:计算-Python-代码的内存和模型显存消耗的小技巧

jiezi

理解 Python 代码的内存耗费是每一个开发人员都必须要解决的问题，这个问题不仅在咱们应用 pandas 读取和解决 CSV 文件的时候十分重要，在咱们应用 GPU 训练的时候还须要布局 GPU 的显存应用。尤其是咱们在白嫖应用 kaggle 和 colab 时显得更为重要。

本篇文章咱们将介绍两个 Python 库 memory_profiler 和 Pytorch-Memory-Utils，这两个库能够帮忙咱们理解内存和显存的耗费。

memory_profiler

pipinstallmemory_profiler#Load its magic function
%load_extmemory_profiler
frommemory_profilerimportprofile

memory_profiler 能够实现以下的工作：

1、查找一行的内存耗费

咱们只须要在代码的后面加上魔法函数 %memit

%memitx = 10+5
#Output
peakmemory: 54.01MiB, increment: 0.27MiB

这里，峰值内存（peak memory）是运行此代码的过程耗费的内存。增量只是因为增加这行代码而须要 / 耗费的内存。同样的逻辑也实用于以下其余的显示。

2、查找函数的内存耗费

在调用函数的行的结尾增加魔法函数。

defaddition():
    a = [1] * (10**1)
    b = [2] * (3*10**2)
    sum = a+b
    returnsum

%memitaddition()
#Output
peakmemory: 36.36MiB, increment: 0.01MiB

3、逐行查找函数的内存耗费

如果须要记录函数中每一行的内存应用，咱们能够应用 @profile 装璜器。然而 @profile 仅实用于在独自模块中定义的函数，因而咱们将首先应用 %%file 创立一个名为 demo.py 的简略模块，其中蕴含咱们的函数 

%%filedemo.py
frommemory_profilerimportprofile

@profile
defaddition():
    a = [1] * (10**1)
    b = [2] * (3*10**2)
    sum = a+b
    returnsum

当初，咱们能够调用该函数 

fromdemoimportaddition

%memitaddition()

#Output
Line#    Mem usage    Increment   Line Contents
================================================
     2     36.4MiB     36.4MiB   @profile
     3                             defaddition():
     4     36.4MiB      0.0MiB       a = [1] * (10**1)
     5   3851.1MiB   3814.7MiB       b = [2] * (3*10**2)
     6   7665.9MiB   3814.8MiB       sum = a+b
     7   7665.9MiB      0.0MiB       returnsum
peakmemory: 7665.88MiB, increment: 7629.52MiB

4、残缺 python 脚本的内存耗费

这个这个办法不能再 notebook 中应用。咱们必须创立 python 脚本并通过命令行运行它。

#create script.py
importtime

@profile
deffunction1():
    n = 100000
    a = [1] *n
    time.sleep(1)
    returna
    
@profile
deffunction2():
    n = 200000
    b = [1] *n
    time.sleep(1)
    returnb

if__name__ == "__main__":
    function1()
    function2()

之后运行脚本并查看 

#On command line
mprofrunscript.py
#To generate plot 
mprofplot

咱们能够看到内存耗费与工夫的关系图

@profile 装璜器没有必要放在函数后面，如果咱们不保留它，咱们不会看到函数级内存耗费，但咱们会看到整个脚本的内存耗费

Pytorch-Memory-Utils

通过 Pytorch-Memory-Utils 工具，咱们在应用显存的代码两头插入检测函数，这样就能够输入在以后行代码时所占用的显存。这个对于咱们计算模型的 GPU 显存占用是十分不便的，通过计算显存占用，咱们才可能最大化训练的 batch size，保障训练的最优速度。

importtorch
importinspect

fromtorchvisionimportmodels
fromgpu_mem_trackimportMemTracker  # 援用显存跟踪代码

device = torch.device('cuda:0')

frame = inspect.currentframe()     
gpu_tracker = MemTracker(frame)      # 创立显存检测对象

gpu_tracker.track()                  # 开始检测
cnn = models.vgg19(pretrained=True).to(device)  # 导入 VGG19 模型并且将数据转到显存中
gpu_tracker.track()

而后能够发现程序运行过程中的显存变动（第一行是载入前的显存，最初一行是载入后的显存）：

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

通过下面的报告，很容易发现一个问题。

首先咱们晓得 VGG19 所有层的权重大小加起来大概是 548M(这个数值来源于 Pytorch 官网提供的 VGG19 权重文件大小)，咱们将下面报告打印的 Tensor-Memory 也都加起来算下来也差不多 551.8Mb。然而，咱们算了两次打印的显存理论占用中：1387.5 – 472.2 = 915.3 MB。

Pytorch 在开始运行程序时须要额定的显存开销，这种额定的显存开销与咱们理论应用的模型权重显存大小无关。这个额定的显存 Pytorch 的开发者也对此进行阐明了，这部分开释后的显存能够用，只不过不在 Nvidia-smi 中显示，所以咱们无需关注。

退出移动版