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

在实验室参与开发了一个评测平台。在使用的时候有部分攻击算法会出现显存溢出的情况。同时随着样本数增加，这种显存占用比会同比增加。而不是和预先设定的一样，仅和设置的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的回答》

但是上述方法是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 "+=" ./

得到结果

[e-mail] /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使用率，大致如图：

但是我想要的是在执行过程中，执行的同时，获取具体的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("内层循环迭代完毕，查看是否正确释放显存")

而这个的输出类似于下图。最好还是添加一个输出提示，因为他如果没有提示做分割的话，其实不是很方便看到底执行到哪里了。

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

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的总量和使用量。但是我实际使用过程中发现，好像并不是非常的好用，具体情况见下图。

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

方法四：找大腿问问

如果debug实在是太难了，也不要一门心思去找自己的问题，找个小伙伴问问。描述给他人的同时你也会更了解问题所在，而且有可能对方一语道破！

最后放个图纪念一下这篇博客的诞生：