GPU Memory Problems in PyTorch（显卡爆炸与利用率不足）

如今研究人工智能，跑深度学习算法，显卡/GPU绝对是第一大门槛，所以不管您是1080Ti还是V100，如果不能发挥出GPU的最大能力，那它可能就是不是显卡而是块普通的砖头了吧。

显卡爆炸
显卡爆炸和内存的使用紧密相连，特别是在代码中对某些变量的不当使用，很有可能内存泄露，从而慢慢得导致显卡OOM（out of memory）。一般来说，计算模型时显存主要是模型参数 + 计算产生的中间变量，细分可以占用分四个部分：

• 模型参数
• 模型计算中间结果
• 反向传播中间结果
• 优化器额外参数

但是如果模型出现显卡内存不足的错误（torch.FatalError: cuda runtime error (2) : out of memory），很难第一时间确定出问题的所在，不过Pycharm的图形化内存分析工具memory profiler可以看到每一部分的时间和函数引用关系，通过这个工具可以确定出问题的某些行代码。说到工具，也有可以通过htop监控下显卡的：

sudo apt-get install htop
htop -d=0.1 #d为更新频率，0.1s更新一次
watch -n 0.1 nvidia-smi #监控内存，0.1s刷新一次

显存比较常见的一些问题有：

• 全局变量累加导致内存泄漏，如loss变量的处理。这里要注意的是，一定要先和GPU detach了，仅仅使用loss的数值，不然添加的实际上是整个计算图。当然除了loss，其他的变量问题基本上都是出现内存泄露的原因呀。

epoch_loss += loss.detach().item()

• 内存垃圾回收。python的变量如果不用了一般不会立即被回收，所以我们可以手动的在每一个epoch里面进行回收。

import gc 
gc.collect()

• 多用np.array，而不是list，特别是DataLoader，而且导入时最好先放到cuda中去。

model = model.cuda()
model.to(device1) #使用序号为1的GPU

• 小心全局变量。
• checkpoint。如果一个模型需要占用的显存实在太大了，那可以使用它分多步进行计算，每次模型都会从上次保存的地方继续训练。
• 爆显存也有可能是忘记设置no_grad（torch.no_grad()是新版本pytorch中volatile的替代）

with torch.no_grad():
   ...

显卡利用率不足
经常有程序跑得非常慢，但是一看占用率才3%，10%，这往往是因为CPU和GPU之间的速度不是很匹配。因为模型的计算基本都是在GPU上的，所以一般问题都出在载入数据的速度太慢了，当载入数据花了很久但模型计算飞快的解决，相对于GPU要等待很久才会工作，这样的占用率自然就不高了。当然也有可能不是数据的问题，而是模型本身就太简单了。

• batch_size变大。这样可以一次性多载入数据到显存中，可以提高它的占用率，并且可以尽量占满GPU的内存。
• Dataloader中的num_workers。这个参数可以多进程的载入数据提高效率，一般可以选用4，8，16等等。但是，这个数量并不是越多越好，因为worker越多，一些进程间的分配和协作+I/O问题反而会拖慢速度。

torch.utils.data.DataLoader(datasets[x],
                            batch_size=batch_size, 
                            shuffle=True,
                            num_workers=8,
                            pin_memory=True)

如果加载一个batch的时间小于模型计算的时间，那么其实不需要多余的worker的，因为GPU不需要额外的等待时间，所以这个时候即便加了worker速度也不会加快。反之如果加载一个batch的时间大于模型计算的时间，那么增加worker之后，可以在前一个worker没就绪前就使用另一个worker的数据来训练。

• pin_memory=True。锁页内存，数据将不在硬盘中存储，省掉了将数据从CPU传入到缓存RAM里面，再给传输到GPU上，利用GPU时就会更快一些。

其他优化显卡的方法

• 使用inplace操作， 比如relu 可以使用 inplace=True。当设置为True时，我们在通过relu()计算时的得到的新值不会占用新的空间而是直接覆盖原来的值，这也就是为什么当inplace参数设置为True时可以节省一部分内存的缘故。

def inplace_relu(m):
    classname = m.__class__.__name__
    if classname.find('ReLU') != -1:
        m.inplace=True

model.apply(inplace_relu)

• del的好习惯，确定某变量不用后，就自己手动清除中间变量。
• 在程序刚开始加这条语句可以提升一点训练速度，没什么额外开销

torch.backends.cudnn.benchmark = True 

• 如果在python内调用pytorch有可能显存和GPU占用不会被自动释放，此时需要加入如下代码来删除一些不需要的变量

 torch.cuda.empty_cache()