深度学习训练中的OOM

1、内存、显存的概念

内存是相对于CPU来说的，而显存是相对于GPU来说的。

2、查看内存和显存使用情况

内存查看命令：top

查看内存使用

可以看到，内存大小是41197352kb，其实就是41G,当前使用了16G

显存查看命令：nvidia-smi

查看显存使用

可以看到显存总量是22G，使用了17G。

一般来说内存都是比显存大的。

3、内存的OOM和显存的OOM

OOM其实分为两种，一种是内存的OOM，一种是显存的OOM。

一般在做模型训练的时候，数据是先读入内存，然后再按batchsize的大小分批次读入显存，做一次训练。

A、内存的OOM

那么，数据在读入内存的时候就可能内存爆掉，比如上万张照片一次性读入内存就会报内存OOM，解决方法就是迭代读入数据。

以keras框架为例，数据迭代读取：

数据迭代读取入内存

在模型训练的时候，使用fit_generator函数就可以了。

模型训练时迭代读取数据

迭代读取数据一般能达到两个效果，一个是降低内存的使用，一个是降低GPU的利用率，但是不会降低显存的使用哈，所以只能解决内存的OOM，不能解决显存的OOM。

B、显存的OOM

首先，来一个显存占用计算公式：

显存占用 = 模型显存占用 + batch_size × 每个样本的显存占用

由公式可以看出，显存占用不是和 batch-size简单正比的关系，尤其是模型自身比较复杂的情况下：比如全连接很大，Embedding 层很大。比如说我曾经做了一个有三个embedding的模型，之后还各种乘，导致网络很复杂。

很显然，内存的OOM，只能通过两种方式来解决：

1、把模型变小：可以通过降低卷积核的个数、卷积层的多少、Embedding的大小、下采样（加池化）、减少全连接层（一般只在最后一层使用）

2、减小batchsize

当然，如果模型实在没有办法变小，batchsize也已经足够小了，还是显存OOM，那就只好使用CPU训练了，因为CPU的内存毕竟比GPU的显存大嘛，只要不超过CPU的内存大小都可以训练，只不过慢一点而已。