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