本文将介绍两个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__
