当前位置: 首页 > 科技观察

Python代码内存和模型内存消耗计算技巧

时间:2023-03-16 00:05:59 科技观察

本文将介绍两个Python库memory_profiler和Pytorch-Memory-Utils,可以帮助我们理解内存和内存消耗。memory_profilerpipinstallmemory_profiler#Loaditsmagicfunction%load_extmemory_profilerfrommemory_profileimportprofilememory_profiler可以完成以下任务:1.求一行的内存消耗,我们只需要加上魔法函数%memit%memitx=10+5#Outputpeakmemory:54.01MiB,increment:0peakvalue27MiB这里,Peakmemory是运行这段代码的进程消耗的内存。增量只是由于添加这行代码而需要/消耗的内存。同样的逻辑适用于下面的其他显示。2.查找函数的内存消耗在调用函数的行的开头添加magicfunction。defaddition():a=[1]*(10**1)b=[2]*(3*10**2)sum=a+breturnsum%memitaddition()#Outputpeakmemory:36.36MiB,increment:0.01MiB3、逐行查找一个函数的内存消耗如果我们需要记录一个函数中每一行的内存使用情况,可以使用@profile装饰器。但是@profile只适用于定义在单独模块中的函数,所以我们首先使用%%file创建一个名为demo.py的简单模块,其中包含我们的函数%%filedemo.pyfrommemory_profileimportprofile@profiledefaddition():a=[1]*(10**1)b=[2]*(3*10**2)sum=a+breturnsum现在,我们可以调用函数fromdemoimportaddition%memitaddition()#OutputLine#MemusageIncrementLineContents======================================================236.4MiB36.4MiB@profile3defaddition():436.4MiB0.0MiBa=[1]*(10**1)53851.1MiB3814.7MiBb=[2]*(3*10**2)67665.9MiB3814.8MiBsum=a+b77665.9MiB0.0MiBreturnspeakmemory:7665.88MiB,increment:7629.52MiB4,完整python脚本的内存消耗这个方法不能在notebook中使用。我们必须创建python脚本并通过命令行运行它。#createscript.pyimporttime@profiledeffunction1():n=100000a=[1]*ntime.sleep(1)return@profiledefunction2():n=200000b=[1]*ntime.sleep(1)returnbif__name__=="__main__":function1()afterfunction2()运行脚本并查看#Oncommandlinemprofrunscript.py#Togenerateplotmprofplot我们可以看到内存消耗与时间函数级内存消耗的图表,但我们将看到整个脚本的内存消耗Pytorch-Memory-Utils通过Pytorch-Memory-Utils这个工具,我们在使用显存的代码中间插入一个检测函数,这样我们就可以输出当前的代码行。的视频内存。这对于我们计算模型的GPU内存使用来说是非常方便的。通过计算内存使用量,我们可以最大化训练的batchsize,保证最优的训练速度。importtorchimportinspectfromtorchvisionimportmodelsfromgpu_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__:line13TotalUsedMemory:472.2Mb+|1*Size:(128,64,3,3)|Memory:0.2949M|+|1*Size:(256,128,3,3)|Memory:1.1796M|+|1*Size:(64,64,3,3)|内存:0.1474M|<类'torch.nn.parameter.Parameter'>+|2*大小:(4096,)|内存:0.0327M|+|1*大小:(512,256,3,3)|内存:4.7185M|+|2*大小:(128,)|内存:0.0010M|+|1*大小:(1000,4096)|内存:16.384M|+|6*大小:(512,)|内存ory:0.0122M|+|1*Size:(64,3,3,3)|Memory:0.0069M|+|1*大小:(4096,25088)|内存:411.04M|+|1*大小:(4096,4096)|内存:67.108M|+|5*Size:(512,512,3,3)|Memory:47.185M|+|2*Size:(64,)|内存:0.0005M|+|3*大小:(256,)|内存:0.0030M|+|1*大小:(128,128,3,3)|内存:0.5898M|+|2*Size:(256,256,3,3)|Memory:4.7185M|+|1*Size:(1000,)|Memory:0.004M|At__main__:line15TotalUsedMemory:1387.5Mb通过上面的报告,很很容易发现问题。首先,我们知道VGG19所有层的权重加起来约为548M(这个值来自Pytorch官方提供的VGG19权重文件的大小)。我们把上面报告打印的Tensor-Memory也加起来,计算也差不多。551.8MB。然而,我们计算了两次打印的实际内存使用量:1387.5–472.2=915.3MB。Pytorch在开始运行程序时需要额外的显存开销。这种额外的显存开销与我们实际使用的模型权重显存大小无关。这个额外显存的开发者Pytorch也对此进行了解释。这部分释放出来的显存是可以使用的,但是在Nvidia-smi中是没有显示的,所以不用关注。