欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

torch代码运行时显存溢出问题

程序员文章站 2022-04-27 07:52:19
...

在实验室参与开发了一个评测平台。在使用的时候有部分攻击算法会出现显存溢出的情况。同时随着样本数增加,这种显存占用比会同比增加。而不是和预先设定的一样,仅和设置的batch_size相关。如此一来,对于一些占用显存本身就较大的算法,当样本数增加时,Docker虚环境肯定会崩溃。下面将给出我的整个搜索+解决问题的过程。

写在最前面的话

这个问题目前已经解决,最终发现报错的原因是,开发那边没有正确的把batch_size传入攻击算法中,导致出现了可能只有1张图像,但是开了一个65倍图像尺寸的空间(实际上3就够了),然后这个空间作为输入传入模型,导致占用显存过多。

在排查问题的过程中,马佬告诉我,其实Pytorch之类的都会有自动回收机制,需要保证的其实是

for循环中的变量,如果是显存上的,尽量不要让他离开for循环范围!
按照GC的原理,是引用计数的,当某个局部变量不存在引用的时候,会自动回收。因此如果for循环内部/外部有引用,都会导致某些中间变量一直被持有。

举个例子:

losses = []
for i in range(233):
    x = Variable(input).to(device)  此时x在GPU上
    output = self.model(x)          此时output也在GPU上
    losses.append(output)           这句话将可能导致存储了output梯度,并由于持有output对象导致他不会在每次for循环后释放
y = x + ...         这句话在for循环外,等于for循环结束的时候,x仍存在未来的引用可能,此时的x不会被回收

可以修改的方式有很多,比如在for循环内部losses.append一句中,可以把output转成cpu上资源。以及将y = 这一句考虑能不能删去。


下面是正文:首先列举全部搜索到的问题,直接跳到最后有几个定位bug小方法。

问题一 记录累计信息时直接使用了输出的Variable

这个问题的发现,是参考了这篇知乎回答《pytorch的坑—loss没写好,显存爆炸》
原贴就问题的描述:

算是动态图的一个坑吧。记录loss信息的时候直接使用了输出的Variable。

for data, label in trainloader:
    out = model(data)
    loss = criterion(out, label)
    loss_sum += loss     # <--- 这里

运行着就发现显存炸了。观察了一下发现随着每个batch显存消耗在不断增大…
参考了别人的代码发现那句loss一般是这样写:

loss_sum += loss.data[0]

这是因为输出的loss的数据类型是Variable。而PyTorch的动态图机制就是通过Variable来构建图。主要是使用Variable计算的时候,会记录下新产生的Variable的运算符号,在反向传播求导的时候进行使用。
如果这里直接将loss加起来,系统会认为这里也是计算图的一部分,也就是说网络会一直延伸变大,那么消耗的显存也就越来越大
总之使用Variable的数据时候要非常小心。不是必要的话尽量使用Tensor来进行计算…

问题二 for循环过程中的迭代变量

参考讨论帖《Tensor to Variable and memory freeing best practices》
在这篇帖子中有提到,Variable和Tensor实际上共用的是一块内存空间。所以在使用了Variable之后,del掉相应的Variable。不会带来明显的内存释放。唯一可能带来一定效果的,是在for循环过程中,如

for i, (x, y) in enumerate(train_loader):
    x = Variable(x)
    y = Variable(y)
    # compute model and update
    del x, y, output 

x和y本身作为train_loader中内容,会占用一块内存,而循环时,会产生一块临时内存。帖子中回复认为,此处可以节省一点点。需要注意的是,还需要额外删去引用到x和y的变量,否则仍然存在占用。

问题三 多次训练,GPU未释放

参考自讨论帖《How can we release GPU memory cache?》
这个帖子中描述的解决办法为,当GPU计算完毕后,把相应的变量和结果转成CPU,然后调用GC,调用torch.cuda.empty_cache()

def wipe_memory(self): # DOES WORK
    self._optimizer_to(torch.device('cpu'))
    del self.optimizer
    gc.collect()
    torch.cuda.empty_cache()

def _optimizer_to(self, device):
    for param in self.optimizer.state.values():
        # Not sure there are any global tensors in the state dict
        if isinstance(param, torch.Tensor):
            param.data = param.data.to(device)
            if param._grad is not None:
                param._grad.data = param._grad.data.to(device)
        elif isinstance(param, dict):
            for subparam in param.values():
                if isinstance(subparam, torch.Tensor):
                    subparam.data = subparam.data.to(device)
                    if subparam._grad is not None:
                        subparam._grad.data = subparam._grad.data.to(device)

问题四 torch.load的坑

参考自知乎回答《PyTorch 有哪些坑/bug? - 知乎用户的回答》
该回答中描述,当你使用:

checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint["state_dict"])

这样load一个 pretrained model 的时候,torch.load() 会默认把load进来的数据放到0卡上,这样每个进程全部会在0卡占用一部分显存。解决的方法也很简单,就是把load进来的数据map到cpu上:

checkpoint = torch.load("checkpoint.pth", map_location=torch.device('cpu'))
model.load_state_dict(checkpoint["state_dict"])

按照马佬的建议,此处如果不想用到cpu的话,也可以map_location=rank。具体的写法参考了《pytorch源码》以及《pytorch 分布式训练 distributed parallel 笔记》

    # 获取GPU的rank号
    gpu = torch.distributed.get_rank(group=group)  # group是可选参数,返回int,执行该脚本的进程的rank
    # 获取了进程号后
    rank = 'cuda:{}'.format(gpu)
    checkpoint = torch.load(args.resume, map_location=rank)

问题五 pretrain weights问题

参考自之乎回答《PyTorch 有哪些坑/bug? - 鲲China的回答》
在做交叉验证的时候,每折初始化模型,由于用到了pretrained weights,这时候显存不会被释放,几折过后显存就爆炸了~,这时候用三行代码就可以解决这个问题

del model
gc.collect()
torch.cuda.empty_cache()

问题六 不做backward,中间变量会保存

参考自《PyTorch 有哪些坑/bug? - hjy666的回答》

torch代码运行时显存溢出问题

但是上述方法是0.4中的解决方法。pytorch0.4到pytrch1.0跨度有点大,variable跟tensor合并成tensor了,不能设置volatile 参数,所以在做evaluation时很容易出现out of memory的问题。所以你需要在最后的loss和predict输出设置

.cpu().detach()

比如说:

total_loss.append(loss.cpu().detach().numpy())
total_finish_loss.append(finish_loss.cpu().detach().numpy())

尝试解决问题

方法一:全局查找字符串

全局查找累计过程,由于主要是+=的问题,所以grep +=试试:

$ grep -rn "+=" ./

得到结果

aaa@qq.com /d/LABOR/SUIBUAA_AIEP (dev_aiep)
$ grep -rn "+=" ./
Binary file ./Datasets/ImageNet/images/ILSVRC2012_val_00000005.JPEG matches
Binary file ./Datasets/ImageNet/images/ILSVRC2012_val_00000006.JPEG matches
Binary file ./Datasets/ImageNet/images/ILSVRC2012_val_00000007.JPEG matches
Binary file ./Datasets/ImageNet/images/ILSVRC2012_val_00000008.JPEG matches
./EvalBox/Defense/anp.py:100:                total += inputs.shape[0]
./EvalBox/Defense/anp.py:101:                correct += (preds == labels).sum().item()
./EvalBox/Defense/eat.py:223:                total += inputs.shape[0]
./EvalBox/Defense/eat.py:224:                correct += (preds == labels).sum().item()
...

但是搜索结果中存在很多的Binary file文件,把所有搜索结果拷贝到sublime中,ctrl + F搜索全部包含"Binary file"字样的搜索行,使用ctrl + shift + K一键删除所有匹配行。

$ grep -rn "+=" ./
./EvalBox/Analysis/grad-cam.py:36:                outputs += [x]
./EvalBox/Analysis/grad-cam.py:134:            cam += w * target[i, :, :]
./EvalBox/Analysis/grand_CAM.py:32:                outputs += [x]
./EvalBox/Analysis/grand_CAM.py:109:            cam += w * target[i, :, :]
./EvalBox/Analysis/Rebust_Defense.py:66:                total += inputs.shape[0]
./EvalBox/Analysis/Rebust_Defense.py:67:                correct += (preds == labels).sum().item()
./EvalBox/Attack/AdvAttack/deepfool.py:105:                loop_i += 1
./EvalBox/Attack/AdvAttack/deepfool.py:137:            loop_i += 1
./EvalBox/Attack/AdvAttack/ead.py:208:                cnt += 1
...

然后再手动筛选掉与该问题无关的行,如上方示例中deepfool中的+=1,这里并不会产生问题一中,无用梯度不释放问题。然后这里我很快就定位到了具体的py文件中,有这么一行

    output = model(xs)

方法二:确定输入输出尺寸

这一步很简单,就是在你觉得不妥的变量上,输出一下他的尺寸看看

print("in line xxx, the var xs 's shape = ", xs.shape)

加一些提示语,然后看看会不会是传入的图像太大了。

我遇到的实际问题就是因为,攻击算法执行过程中,用于做扰动处理的预空间维度太高。按照马佬的测试,1张3 * 244 * 244的ImageNet图像,在VGG模型上执行预测,约占用显存1.6G。而我传入的是80 * 3 * 375 * 500的输入,所以显存爆炸。改为3 * 3 * 375 * 500之后,显存就可以正常供给了。

方法三:如何查看实时的GPU使用率

这个也是debug过程中很苦恼的东西,想知道是不是在某一步的时候,传到显存上的东西太多了,但是又不方便单步调试。

使用指令nvidia-smi可以看到当前的GPU使用率,大致如图:
torch代码运行时显存溢出问题

但是我想要的是在执行过程中,执行的同时,获取具体的GPU使用情况。这里我参考了这篇博客《使用python中的GPUtil库从NVIDA GPU获取GPU状态》
这里面用到了一个第三方库叫GPUtil,执行pip install gputil即可完成下载。然后我封装了一个函数:

   def get_gpu_info(self, text = ""):
      print("当前行为为:", text)
      GPUtil.showUtilization()
   def predict(self, xs, model):
       var_xs = Variable(xs.to(device))
       self.get_gpu_info("将xs传入GPU")
       for i in range(100):
           for j in range(200):
               output = model(var_xs)
               self.get_gpu_info("执行一次预测过程")
               some work there ...
           self.get_gpu_info("内层循环迭代完毕,查看是否正确释放显存")

而这个的输出类似于下图。最好还是添加一个输出提示,因为他如果没有提示做分割的话,其实不是很方便看到底执行到哪里了。
torch代码运行时显存溢出问题

在这篇参考博客中,我看到有这么一段代码

import GPUtil
import time
while True:
    Gpus = GPUtil.getGPUs()
    for gpu in Gpus:
        print('GPU总量', gpu.memoryTotal)
        print('GPU使用量', gpu.memortUsed)
    time.sleep(5)

他这里的意思是不停的输出GPU的总量和使用量。但是我实际使用过程中发现,好像并不是非常的好用,具体情况见下图。
torch代码运行时显存溢出问题

可以看到右边是我的一个实测结果,虽然我中间改变了GPU的使用情况,但是输出的值基本没变。不知道是更新不够快还是如何。我一开始以为是更新不够快,但是我发现哪怕程序一开始,他都可能会显示已经占用了部分的显存资源。所以我就改用了上面的那个GPUtil.showUtilization()

方法四:找大腿问问

如果debug实在是太难了,也不要一门心思去找自己的问题,找个小伙伴问问。描述给他人的同时你也会更了解问题所在,而且有可能对方一语道破!
torch代码运行时显存溢出问题

最后放个图纪念一下这篇博客的诞生:
torch代码运行时显存溢出问题

相关标签: pytorch oom