如今研究人工智能,跑深度学习算法,显卡/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()